Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f6a6fa6

Browse files
dstansbyclairefi0ksundenQuLogic
authored
Add new num_arrows option to streamplot (#27617)
* Add new narrows option to streamplot Fix test image numarrows > narrows * Add multiple arrows example Fix example * Tidy up alt text in what's new example * narrows > n_arrows * Fix new streamplot type stub Co-authored-by: Kyle Sunden <[email protected]> * Rename n_arrows to num_arrows * Fix handling of stream line width with multiple arrows --------- Co-authored-by: clairefio <[email protected]> Co-authored-by: Kyle Sunden <[email protected]> Co-authored-by: Elliott Sales de Andrade <[email protected]>
1 parent 2fa9151 commit f6a6fa6

File tree

10 files changed

+73
-3124
lines changed

10 files changed

+73
-3124
lines changed

‎doc/users/next_whats_new/streamplot_multiple_arrows.rst

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
Multiple arrows on a streamline
2+
-------------------------------
3+
4+
A new ``num_arrows`` argument has been added to `~matplotlib.axes.Axes.streamplot` that
5+
allows more than one arrow to be added to each streamline:
6+
7+
.. plot::
8+
:include-source: true
9+
:alt: One chart showing a streamplot. Each streamline has three arrows.
10+
11+
import matplotlib.pyplot as plt
12+
import numpy as np
13+
14+
w = 3
15+
Y, X = np.mgrid[-w:w:100j, -w:w:100j]
16+
U = -1 - X**2 + Y
17+
V = 1 + X - Y**2
18+
19+
fig, ax = plt.subplots()
20+
ax.streamplot(X, Y, U, V, num_arrows=3)
21+
22+
plt.show()

‎galleries/examples/images_contours_and_fields/plot_streamplot.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
V = 1 + X - Y**2
2424
speed = np.sqrt(U**2 + V**2)
2525

26-
fig, axs = plt.subplots(3, 2, figsize=(7, 9), height_ratios=[1, 1, 2])
26+
fig, axs = plt.subplots(4, 2, figsize=(7, 12), height_ratios=[1, 1, 1, 2])
2727
axs = axs.flat
2828

2929
# Varying density along a streamline
@@ -37,7 +37,7 @@
3737

3838
# Varying line width along a streamline
3939
lw = 5*speed / speed.max()
40-
axs[2].streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw)
40+
axs[2].streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw, num_arrows=5)
4141
axs[2].set_title('Varying Line Width')
4242

4343
# Controlling the starting points of the streamlines
@@ -52,21 +52,27 @@
5252
axs[3].plot(seed_points[0], seed_points[1], 'bo')
5353
axs[3].set(xlim=(-w, w), ylim=(-w, w))
5454

55+
# Adding more than one arrow to each streamline
56+
axs[4].streamplot(X, Y, U, V, num_arrows=3)
57+
axs[4].set_title('Multiple arrows')
58+
59+
axs[5].axis("off")
60+
5561
# Create a mask
5662
mask = np.zeros(U.shape, dtype=bool)
5763
mask[40:60, 40:60] = True
5864
U[:20, :20] = np.nan
5965
U = np.ma.array(U, mask=mask)
6066

61-
axs[4].streamplot(X, Y, U, V, color='r')
62-
axs[4].set_title('Streamplot with Masking')
67+
axs[6].streamplot(X, Y, U, V, color='r')
68+
axs[6].set_title('Streamplot with Masking')
6369

64-
axs[4].imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray',
70+
axs[6].imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray',
6571
aspect='auto')
66-
axs[4].set_aspect('equal')
72+
axs[6].set_aspect('equal')
6773

68-
axs[5].streamplot(X, Y, U, V, broken_streamlines=False)
69-
axs[5].set_title('Streamplot with unbroken streamlines')
74+
axs[7].streamplot(X, Y, U, V, broken_streamlines=False)
75+
axs[7].set_title('Streamplot with unbroken streamlines')
7076

7177
plt.tight_layout()
7278
plt.show()

‎lib/matplotlib/pyplot.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4125,6 +4125,7 @@ def streamplot(
41254125
integration_direction="both",
41264126
broken_streamlines=True,
41274127
*,
4128+
num_arrows=1,
41284129
data=None,
41294130
):
41304131
__ret = gca().streamplot(
@@ -4146,6 +4147,7 @@ def streamplot(
41464147
maxlength=maxlength,
41474148
integration_direction=integration_direction,
41484149
broken_streamlines=broken_streamlines,
4150+
num_arrows=num_arrows,
41494151
**({"data": data} if data is not None else {}),
41504152
)
41514153
sci(__ret.lines)

‎lib/matplotlib/streamplot.py

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
1919
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
2020
minlength=0.1, transform=None, zorder=None, start_points=None,
2121
maxlength=4.0, integration_direction='both',
22-
broken_streamlines=True):
22+
broken_streamlines=True, *, num_arrows=1):
2323
"""
2424
Draw streamlines of a vector flow.
2525
@@ -73,6 +73,11 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
7373
If False, forces streamlines to continue until they
7474
leave the plot domain. If True, they may be terminated if they
7575
come too close to another streamline.
76+
num_arrows : int
77+
Number of arrows per streamline. The arrows are spaced equally along the steps
78+
each streamline takes. Note that this can be different to being spaced equally
79+
along the distance of the streamline.
80+
7681
7782
Returns
7883
-------
@@ -92,6 +97,9 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
9297
mask = StreamMask(density)
9398
dmap = DomainMap(grid, mask)
9499

100+
if num_arrows < 0:
101+
raise ValueError(f"The value of num_arrows must be >= 0, got {num_arrows=}")
102+
95103
if zorder is None:
96104
zorder = mlines.Line2D.zorder
97105

@@ -206,25 +214,31 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
206214
points = np.transpose([tx, ty])
207215
streamlines.append(points)
208216

209-
# Add arrows halfway along each trajectory.
217+
# Distance along streamline
210218
s = np.cumsum(np.hypot(np.diff(tx), np.diff(ty)))
211-
n = np.searchsorted(s, s[-1] / 2.)
212-
arrow_tail = (tx[n], ty[n])
213-
arrow_head = (np.mean(tx[n:n + 2]), np.mean(ty[n:n + 2]))
214-
215219
if isinstance(linewidth, np.ndarray):
216220
line_widths = interpgrid(linewidth, tgx, tgy)[:-1]
217221
line_kw['linewidth'].extend(line_widths)
218-
arrow_kw['linewidth'] = line_widths[n]
219-
220222
if use_multicolor_lines:
221223
color_values = interpgrid(color, tgx, tgy)[:-1]
222224
line_colors.append(color_values)
223-
arrow_kw['color'] = cmap(norm(color_values[n]))
224225

225-
p = patches.FancyArrowPatch(
226-
arrow_tail, arrow_head, transform=transform, **arrow_kw)
227-
arrows.append(p)
226+
# Add arrows along each trajectory.
227+
for x in range(1, num_arrows+1):
228+
# Get index of distance along streamline to place arrow
229+
idx = np.searchsorted(s, s[-1] * (x/(num_arrows+1)))
230+
arrow_tail = (tx[idx], ty[idx])
231+
arrow_head = (np.mean(tx[idx:idx + 2]), np.mean(ty[idx:idx + 2]))
232+
233+
if isinstance(linewidth, np.ndarray):
234+
arrow_kw['linewidth'] = line_widths[idx]
235+
236+
if use_multicolor_lines:
237+
arrow_kw['color'] = cmap(norm(color_values[idx]))
238+
239+
p = patches.FancyArrowPatch(
240+
arrow_tail, arrow_head, transform=transform, **arrow_kw)
241+
arrows.append(p)
228242

229243
lc = mcollections.LineCollection(
230244
streamlines, transform=transform, **line_kw)

‎lib/matplotlib/streamplot.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ def streamplot(
2828
maxlength: float = ...,
2929
integration_direction: Literal["forward", "backward", "both"] = ...,
3030
broken_streamlines: bool = ...,
31+
*,
32+
num_arrows: int = ...,
3133
) -> StreamplotSet: ...
3234

3335
class StreamplotSet:

‎lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_linewidth.pdf

-21.3 KB
Binary file not shown.

‎lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_linewidth.png

1.93 KB
Loading

0 commit comments

Comments
 (0)