matplotlib · ksunden · Nov 20, 2024 · Jun 24, 2020 · Jan 9, 2024 · Jan 10, 2024
diff --git a/doc/users/next_whats_new/streamplot_multiple_arrows.rst b/doc/users/next_whats_new/streamplot_multiple_arrows.rst
@@ -0,0 +1,22 @@
+Multiple arrows on a streamline
+-------------------------------
+
+A new ``num_arrows`` argument has been added to `~matplotlib.axes.Axes.streamplot` that
+allows more than one arrow to be added to each streamline:
+
+.. plot::
+    :include-source: true
+    :alt: One chart showing a streamplot. Each streamline has three arrows.
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    w = 3
+    Y, X = np.mgrid[-w:w:100j, -w:w:100j]
+    U = -1 - X**2 + Y
+    V = 1 + X - Y**2
+
+    fig, ax = plt.subplots()
+    ax.streamplot(X, Y, U, V, num_arrows=3)
+
+    plt.show()
diff --git a/galleries/examples/images_contours_and_fields/plot_streamplot.py b/galleries/examples/images_contours_and_fields/plot_streamplot.py
@@ -23,7 +23,7 @@
 V = 1 + X - Y**2
 speed = np.sqrt(U**2 + V**2)
 
-fig, axs = plt.subplots(3, 2, figsize=(7, 9), height_ratios=[1, 1, 2])
+fig, axs = plt.subplots(4, 2, figsize=(7, 12), height_ratios=[1, 1, 1, 2])
 axs = axs.flat
 
 #  Varying density along a streamline
@@ -37,7 +37,7 @@
 
 #  Varying line width along a streamline
 lw = 5*speed / speed.max()
-axs[2].streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw)
+axs[2].streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw, num_arrows=5)
 axs[2].set_title('Varying Line Width')
 
 # Controlling the starting points of the streamlines
@@ -52,21 +52,27 @@
 axs[3].plot(seed_points[0], seed_points[1], 'bo')
 axs[3].set(xlim=(-w, w), ylim=(-w, w))
 
+# Adding more than one arrow to each streamline
+axs[4].streamplot(X, Y, U, V, num_arrows=3)
+axs[4].set_title('Multiple arrows')
+
+axs[5].axis("off")
+
 # Create a mask
 mask = np.zeros(U.shape, dtype=bool)
 mask[40:60, 40:60] = True
 U[:20, :20] = np.nan
 U = np.ma.array(U, mask=mask)
 
-axs[4].streamplot(X, Y, U, V, color='r')
-axs[4].set_title('Streamplot with Masking')
+axs[6].streamplot(X, Y, U, V, color='r')
+axs[6].set_title('Streamplot with Masking')
 
-axs[4].imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray',
+axs[6].imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray',
               aspect='auto')
-axs[4].set_aspect('equal')
+axs[6].set_aspect('equal')
 
-axs[5].streamplot(X, Y, U, V, broken_streamlines=False)
-axs[5].set_title('Streamplot with unbroken streamlines')
+axs[7].streamplot(X, Y, U, V, broken_streamlines=False)
+axs[7].set_title('Streamplot with unbroken streamlines')
 
 plt.tight_layout()
 plt.show()

diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -4125,6 +4125,7 @@ def streamplot(
     integration_direction="both",
     broken_streamlines=True,
     *,
+    num_arrows=1,
     data=None,
 ):
     __ret = gca().streamplot(
@@ -4146,6 +4147,7 @@ def streamplot(
         maxlength=maxlength,
         integration_direction=integration_direction,
         broken_streamlines=broken_streamlines,
+        num_arrows=num_arrows,
         **({"data": data} if data is not None else {}),
     )
     sci(__ret.lines)

diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -19,7 +19,7 @@
                cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
                minlength=0.1, transform=None, zorder=None, start_points=None,
                maxlength=4.0, integration_direction='both',
-               broken_streamlines=True):
+               broken_streamlines=True, *, num_arrows=1):
     """
     Draw streamlines of a vector flow.
 
@@ -73,6 +73,11 @@
         If False, forces streamlines to continue until they
         leave the plot domain.  If True, they may be terminated if they
         come too close to another streamline.
+    num_arrows : int
+        Number of arrows per streamline. The arrows are spaced equally along the steps
+        each streamline takes. Note that this can be different to being spaced equally
+        along the distance of the streamline.
+
 
     Returns
     -------
@@ -92,6 +97,9 @@
     mask = StreamMask(density)
     dmap = DomainMap(grid, mask)
 
+    if num_arrows < 0:
+        raise ValueError(f"The value of num_arrows must be >= 0, got {num_arrows=}")
+
     if zorder is None:
         zorder = mlines.Line2D.zorder
 
@@ -206,25 +214,31 @@
             points = np.transpose([tx, ty])
             streamlines.append(points)
 
-        # Add arrows halfway along each trajectory.
+        # Distance along streamline
         s = np.cumsum(np.hypot(np.diff(tx), np.diff(ty)))
-        n = np.searchsorted(s, s[-1] / 2.)
-        arrow_tail = (tx[n], ty[n])
-        arrow_head = (np.mean(tx[n:n + 2]), np.mean(ty[n:n + 2]))
-
         if isinstance(linewidth, np.ndarray):
             line_widths = interpgrid(linewidth, tgx, tgy)[:-1]
             line_kw['linewidth'].extend(line_widths)
-            arrow_kw['linewidth'] = line_widths[n]
-
         if use_multicolor_lines:
             color_values = interpgrid(color, tgx, tgy)[:-1]
             line_colors.append(color_values)
-            arrow_kw['color'] = cmap(norm(color_values[n]))
 
-        p = patches.FancyArrowPatch(
-            arrow_tail, arrow_head, transform=transform, **arrow_kw)
-        arrows.append(p)
+        # Add arrows along each trajectory.
+        for x in range(1, num_arrows+1):
+            # Get index of distance along streamline to place arrow
+            idx = np.searchsorted(s, s[-1] * (x/(num_arrows+1)))
+            arrow_tail = (tx[idx], ty[idx])
+            arrow_head = (np.mean(tx[idx:idx + 2]), np.mean(ty[idx:idx + 2]))
+
+            if isinstance(linewidth, np.ndarray):
+                arrow_kw['linewidth'] = line_widths[idx]
+
+            if use_multicolor_lines:
+                arrow_kw['color'] = cmap(norm(color_values[idx]))
+
+            p = patches.FancyArrowPatch(
+                arrow_tail, arrow_head, transform=transform, **arrow_kw)
+            arrows.append(p)
 
     lc = mcollections.LineCollection(
         streamlines, transform=transform, **line_kw)

diff --git a/lib/matplotlib/streamplot.pyi b/lib/matplotlib/streamplot.pyi
@@ -28,6 +28,8 @@ def streamplot(
     maxlength: float = ...,
     integration_direction: Literal["forward", "backward", "both"] = ...,
     broken_streamlines: bool = ...,
+    *,
+    num_arrows: int = ...,
 ) -> StreamplotSet: ...
 
 class StreamplotSet:

diff --git a/lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_linewidth.pdf b/lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_linewidth.pdf
diff --git a/lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_linewidth.png b/lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_linewidth.png