fill_between extended to 3D

scottshambaugh · scottshambaugh · commit d1607afa46d8 · 2024-06-13T08:03:13.000-06:00
fill_between in plot types

Apply suggestions from code review

Co-authored-by: Tim Hoffmann &lt;2836374+timhoffm@users.noreply.github.com&gt;

fill_between single_polygon flag
diff --git a/doc/api/toolkits/mplot3d/axes3d.rst b/doc/api/toolkits/mplot3d/axes3d.rst
@@ -30,6 +30,7 @@ Plotting
    plot_surface
    plot_wireframe
    plot_trisurf
+   fill_between
 
    clabel
    contour
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,25 @@
+Fill between 3D lines
+---------------------
+
+The new method `.Axes3D.fill_between` allows to fill the surface between two
+3D lines with polygons.
+
+.. plot::
+    :include-source:
+    :alt: Example of 3D fill_between
+
+    N = 50
+    theta = np.linspace(0, 2*np.pi, N)
+
+    x1 = np.cos(theta)
+    y1 = np.sin(theta)
+    z1 = 0.1 * np.sin(6 * theta)
+
+    x2 = 0.6 * np.cos(theta)
+    y2 = 0.6 * np.sin(theta)
+    z2 = 2  # Note that scalar values work in addition to length N arrays
+
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    ax.fill_between(x1, y1, z1, x2, y2, z2,
+                    alpha=0.5, edgecolor='k')
diff --git a/galleries/examples/mplot3d/fillbetween3d.py b/galleries/examples/mplot3d/fillbetween3d.py
@@ -0,0 +1,28 @@
+"""
+=====================
+Fill between 3D lines
+=====================
+
+Demonstrate how to fill the space between 3D lines with surfaces. Here we
+create a sort of "lampshade" shape.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+N = 50
+theta = np.linspace(0, 2*np.pi, N)
+
+x1 = np.cos(theta)
+y1 = np.sin(theta)
+z1 = 0.1 * np.sin(6 * theta)
+
+x2 = 0.6 * np.cos(theta)
+y2 = 0.6 * np.sin(theta)
+z2 = 2  # Note that scalar values work in addition to length N arrays
+
+fig = plt.figure()
+ax = fig.add_subplot(projection='3d')
+ax.fill_between(x1, y1, z1, x2, y2, z2, alpha=0.5, edgecolor='k')
+
+plt.show()
diff --git a/galleries/examples/mplot3d/fillunder3d.py b/galleries/examples/mplot3d/fillunder3d.py
@@ -0,0 +1,37 @@
+"""
+=========================
+Fill under 3D line graphs
+=========================
+
+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
+    # single_polygon=True is set for faster and cleaner rendering for this
+    # shape on a single plane.
+    ax.fill_between(x, l, l**x * np.exp(-l) / gamma(x + 1),
+                    x, l, 0,
+                    single_polygon=True,
+                    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.
 """
 
-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 = [-3, -2, -1]
 
-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*z, y*z, np.full_like(x, z))))
+verts = np.array(verts)
 
 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, alpha=.7)
+ax.add_collection3d(poly)
+ax.auto_scale_xyz(verts[:, :, 0], verts[:, :, 1], verts[:, :, 2])
+ax.set_aspect('equalxy')
 
 plt.show()
diff --git a/galleries/plot_types/3D/fill_between3d_simple.py b/galleries/plot_types/3D/fill_between3d_simple.py
@@ -0,0 +1,33 @@
+"""
+====================================
+fill_between(x1, y1, z1, x2, y2, z2)
+====================================
+
+See `~mpl_toolkits.mplot3d.axes3d.Axes3D.fill_between`.
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+plt.style.use('_mpl-gallery')
+
+# Make data
+n = 50
+theta = np.linspace(0, 2*np.pi, n)
+x1 = np.cos(theta)
+y1 = np.sin(theta)
+z1 = np.linspace(0, 1, n)
+x2 = np.cos(theta + np.pi)
+y2 = np.sin(theta + np.pi)
+z2 = z1
+
+# Plot
+fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
+ax.fill_between(x1, y1, z1, x2, y2, z2, alpha=0.5)
+ax.plot(x1, y1, z1, linewidth=2, color='C0')
+ax.plot(x2, y2, z2, linewidth=2, color='C0')
+
+ax.set(xticklabels=[],
+       yticklabels=[],
+       zticklabels=[])
+
+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 3D lines
+=====================
+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
@@ -1957,6 +1957,99 @@ def plot(self, xs, ys, *args, zdir='z', **kwargs):
 
     plot3D = plot
 
+    def fill_between(self, x1, y1, z1, x2, y2, z2, *,
+                     where=None, single_polygon=False, **kwargs):
+        """
+        Fill the area between two 3D curves.
+
+        The curves are defined by the points (*x1*, *y1*, *z1*) and
+        (*x2*, *y2*, *z2*). 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
+        ----------
+        x1, y1, z1 : float or 1D array-like
+            x, y, and z  coordinates of vertices for 1st line.
+
+        x2, y2, z2 : 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.
+
+        single_polygon : bool
+            If True, then the 1st and 2nd lines are connected to form a single
+            polygon, which may be faster and render more cleanly for simple
+            shapes (eg, for filling between two lines that lie on a single
+            plane). If False (default), then a separate polygon is created for
+            each pair of subsequent points in the 1st and 2nd lines. This is
+            more flexible, but may be slower and result in rendering artifacts.
+
+        **kwargs
+            All other keyword arguments are passed on to `.Poly3DCollection`.
+
+        Returns
+        -------
+        `.Poly3DCollection`
+            A `.Poly3DCollection` containing the plotted polygons.
+
+        """
+        had_data = self.has_data()
+        x1, y1, z1, x2, y2, z2 = cbook._broadcast_with_masks(x1, y1, z1, x2, y2, z2)
+
+        if where is None:
+            where = True
+        else:
+            where = np.asarray(where, dtype=bool)
+            if where.size != x1.size:
+                raise ValueError(f"where size ({where.size}) does not match "
+                                 f"size ({x1.size})")
+        where = where & ~np.isnan(x1)  # NaNs were broadcast in _broadcast_with_masks
+
+        if single_polygon:
+            poly = []
+            for i in range(len(x1)):
+                if where[i]:
+                    poly.append((x1[i], y1[i], z1[i]))
+            for i in range(len(x2) - 1, -1, -1):
+                if where[i]:
+                    poly.append((x2[i], y2[i], z2[i]))
+            polys = [poly]
+
+        else:
+            polys = []
+            for idx0, idx1 in cbook.contiguous_regions(where):
+                x1slice = x1[idx0:idx1]
+                y1slice = y1[idx0:idx1]
+                z1slice = z1[idx0:idx1]
+                x2slice = x2[idx0:idx1]
+                y2slice = y2[idx0:idx1]
+                z2slice = z2[idx0:idx1]
+
+                if not len(x1slice):
+                    continue
+
+                for i in range(len(x1slice) - 1):
+                    poly = [(x1slice[i], y1slice[i], z1slice[i]),
+                            (x1slice[i+1], y1slice[i+1], z1slice[i+1]),
+                            (x2slice[i+1], y2slice[i+1], z2slice[i+1]),
+                            (x2slice[i], y2slice[i], z2slice[i])]
+                    polys.append(poly)
+
+        polyc = art3d.Poly3DCollection(polys, **kwargs)
+        self.add_collection(polyc)
+
+        self.auto_scale_xyz([x1, x2], [y1, y2], [z1, z2], 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_multi_polygon.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/fill_between_multi_polygon.png
diff --git a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/fill_between_single_polygon.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/fill_between_single_polygon.png
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -593,6 +593,44 @@ def test_plot_3d_from_2d():
     ax.plot(xs, ys, zs=0, zdir='y')
 
 
+@mpl3d_image_comparison(['fill_between_multi_polygon.png'], style='mpl20')
+def test_fill_between_multi_polygon():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    theta = np.linspace(0, 2*np.pi, 50)
+
+    x1 = np.cos(theta)
+    y1 = np.sin(theta)
+    z1 = 0.1 * np.sin(6 * theta)
+
+    x2 = 0.6 * np.cos(theta)
+    y2 = 0.6 * np.sin(theta)
+    z2 = 2
+
+    where = (theta < np.pi/2) | (theta > 3*np.pi/2)
+
+    ax.fill_between(x1, y1, z1, x2, y2, z2,
+                    where=where, alpha=0.5, edgecolor='k')
+
+
+@mpl3d_image_comparison(['fill_between_single_polygon.png'], style='mpl20')
+def test_fill_between_single_polygon():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    theta = np.linspace(0, 2*np.pi, 50)
+
+    x1 = x2 = theta
+    y1 = y2 = 0
+    z1 = np.cos(theta)
+    z2 = 0
+
+    ax.fill_between(x1, y1, z1, x2, y2, z2,
+                    single_polygon=True,
+                    alpha=0.5, edgecolor='k')
+
+
 @mpl3d_image_comparison(['surface3d.png'], style='mpl20')
 def test_surface3d():
     # Remove this line when this test image is regenerated.
