Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit f6a6fa6

Browse filesBrowse 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 <git@ksunden.space> * Rename n_arrows to num_arrows * Fix handling of stream line width with multiple arrows --------- Co-authored-by: clairefio <clairefiorino@icloud.com> Co-authored-by: Kyle Sunden <git@ksunden.space> Co-authored-by: Elliott Sales de Andrade <quantum.analyst@gmail.com>
1 parent 2fa9151 commit f6a6fa6
Copy full SHA for f6a6fa6

File tree

Expand file treeCollapse file tree

10 files changed

+73
-3124
lines changed
Filter options
Expand file treeCollapse file tree

10 files changed

+73
-3124
lines changed

‎doc/users/next_whats_new/streamplot_multiple_arrows.rst

Copy file name to clipboard
+22Lines 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

Copy file name to clipboardExpand all lines: galleries/examples/images_contours_and_fields/plot_streamplot.py
+14-8Lines 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

Copy file name to clipboardExpand all lines: lib/matplotlib/pyplot.py
+2Lines 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

Copy file name to clipboardExpand all lines: lib/matplotlib/streamplot.py
+26-12Lines 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

Copy file name to clipboardExpand all lines: lib/matplotlib/streamplot.pyi
+2Lines 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

Copy file name to clipboard
-21.3 KB
Binary file not shown.

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

Copy file name to clipboard
1.93 KB
Loading

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.