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

Skip to content

Commit 604936f

Browse files
eytanadlereytanadler
committed
Add arguments to streamplot to control integration max step and error
1 parent bb3d5e0 commit 604936f

File tree

7 files changed

+201
-10
lines changed

7 files changed

+201
-10
lines changed

‎doc/users/next_whats_new/streamplot_integration_control.rst

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
Streamplot integration control
2+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
3+
4+
Two new options have been added to the `~.axes.Axes.streamplot` function that
5+
give the user better control of the streamline integration. The first is called
6+
``integration_max_step_scale`` and multiplies the default max step computed by the
7+
integrator. The second is called ``integration_max_error_scale`` and multiplies the
8+
default max error set by the integrator. Values for these parameters between
9+
zero and one reduce (tighten) the max step or error to improve streamline
10+
accuracy by performing more computation. Values greater than one increase
11+
(loosen) the max step or error to reduce computation time at the cost of lower
12+
streamline accuracy.
13+
14+
The integrator defaults are both hand-tuned values and may not be applicable to
15+
all cases, so this allows customizing the behavior to specific use cases.
16+
Modifying only ``integration_max_step_scale`` has proved effective, but it may be useful
17+
to control the error as well.

‎galleries/examples/images_contours_and_fields/plot_streamplot.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
* Unbroken streamlines even when exceeding the limit of lines within a single
1515
grid cell.
1616
"""
17+
import time
18+
1719
import matplotlib.pyplot as plt
1820
import numpy as np
1921

@@ -74,6 +76,72 @@
7476
axs[7].streamplot(X, Y, U, V, broken_streamlines=False)
7577
axs[7].set_title('Streamplot with unbroken streamlines')
7678

79+
plt.tight_layout()
80+
plt.show()
81+
82+
# %%
83+
# Streamline computation
84+
# ----------------------
85+
#
86+
# The streamlines are computed by integrating along the provided vector field
87+
# from the seed points, which are either automatically generated or manually
88+
# specified. The accuracy and smoothness of the streamlines can be adjusted using
89+
# the ``integration_max_step_scale`` and ``integration_max_error_scale`` optional
90+
# parameters. See the `~.axes.Axes.streamplot` function documentation for more
91+
# details.
92+
#
93+
# This example shows how adjusting the maximum allowed step size and error for
94+
# the integrator changes the appearance of the streamline. The differences can
95+
# be subtle, but can be observed particularly where the streamlines have
96+
# high curvature.
97+
98+
# Linear potential flow over a lifting cylinder
99+
n = 50
100+
x, y = np.meshgrid(np.linspace(-2, 2, n), np.linspace(-3, 3, n))
101+
th = np.arctan2(y, x)
102+
r = np.sqrt(x**2 + y**2)
103+
vr = -np.cos(th) / r**2
104+
vt = -np.sin(th) / r**2 - 1 / r
105+
vx = vr * np.cos(th) - vt * np.sin(th) + 1.0
106+
vy = vr * np.sin(th) + vt * np.cos(th)
107+
108+
# Seed points
109+
n_seed = 50
110+
seed_pts = np.column_stack((np.full(n_seed, -1.75), np.linspace(-2, 2, n_seed)))
111+
112+
_, axs = plt.subplots(3, 1, figsize=(6, 14))
113+
for ax, max_val in zip(axs, [0.05, 1, 5]):
114+
t_start = time.time()
115+
ax.streamplot(
116+
x,
117+
y,
118+
vx,
119+
vy,
120+
start_points=seed_pts,
121+
broken_streamlines=False,
122+
arrowsize=1e-10,
123+
linewidth=0.6,
124+
color="k",
125+
integration_max_step_scale=max_val,
126+
integration_max_error_scale=max_val,
127+
)
128+
129+
text = f"integration_max_step_scale: {max_val}\n" \
130+
f"integration_max_error_scale: {max_val}\n" \
131+
f"streamplot time: {time.time() - t_start:.2f} sec"
132+
if max_val == 1:
133+
text += "\n(default)"
134+
ax.text(0.0, 0.0, text, ha="center", va="center")
135+
136+
# Draw the cylinder
137+
th_circ = np.linspace(0, 2 * np.pi, 100)
138+
for ax in axs:
139+
ax.fill(np.cos(th_circ), np.sin(th_circ), color="w", ec="k")
140+
141+
ax.set_aspect("equal")
142+
ax.set_ylim(-1.5, 1.5)
143+
ax.axis("off")
144+
77145
plt.tight_layout()
78146
plt.show()
79147
# %%

‎lib/matplotlib/pyplot.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4128,6 +4128,8 @@ def streamplot(
41284128
integration_direction="both",
41294129
broken_streamlines=True,
41304130
*,
4131+
integration_max_step_scale=1.0,
4132+
integration_max_error_scale=1.0,
41314133
num_arrows=1,
41324134
data=None,
41334135
):
@@ -4150,6 +4152,8 @@ def streamplot(
41504152
maxlength=maxlength,
41514153
integration_direction=integration_direction,
41524154
broken_streamlines=broken_streamlines,
4155+
integration_max_step_scale=integration_max_step_scale,
4156+
integration_max_error_scale=integration_max_error_scale,
41534157
num_arrows=num_arrows,
41544158
**({"data": data} if data is not None else {}),
41554159
)

‎lib/matplotlib/streamplot.py

Lines changed: 52 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ 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, *, num_arrows=1):
22+
broken_streamlines=True, *, integration_max_step_scale=1.0,
23+
integration_max_error_scale=1.0, num_arrows=1):
2324
"""
2425
Draw streamlines of a vector flow.
2526
@@ -73,6 +74,26 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
7374
If False, forces streamlines to continue until they
7475
leave the plot domain. If True, they may be terminated if they
7576
come too close to another streamline.
77+
integration_max_step_scale : float, default: 1.0
78+
Multiplier on the maximum allowable step in the streamline integration routine.
79+
A value between zero and one results in a max integration step smaller than
80+
the default max step, resulting in more accurate streamlines at the cost
81+
of greater computation time. A value greater than one results in a max
82+
integration step larger than the default value, reducing computation time
83+
at the cost of less accurate streamlines. Must be greater than zero.
84+
85+
.. versionadded:: 3.11
86+
87+
integration_max_error_scale : float, default: 1.0
88+
Multiplier on the maximum allowable error in the streamline integration routine.
89+
A value between zero and one results in a tighter max integration error than
90+
the default max error, resulting in more accurate streamlines at the cost
91+
of greater computation time. A value greater than one results in a looser
92+
max integration error than the default value, reducing computation time at
93+
the cost of less accurate streamlines. Must be greater than zero.
94+
95+
.. versionadded:: 3.11
96+
7697
num_arrows : int
7798
Number of arrows per streamline. The arrows are spaced equally along the steps
7899
each streamline takes. Note that this can be different to being spaced equally
@@ -97,6 +118,18 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
97118
mask = StreamMask(density)
98119
dmap = DomainMap(grid, mask)
99120

121+
if integration_max_step_scale <= 0.0:
122+
raise ValueError(
123+
"The value of integration_max_step_scale must be > 0, " +
124+
f"got {integration_max_step_scale}"
125+
)
126+
127+
if integration_max_error_scale <= 0.0:
128+
raise ValueError(
129+
"The value of integration_max_error_scale must be > 0, " +
130+
f"got {integration_max_error_scale}"
131+
)
132+
100133
if num_arrows < 0:
101134
raise ValueError(f"The value of num_arrows must be >= 0, got {num_arrows=}")
102135

@@ -160,7 +193,9 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
160193
for xm, ym in _gen_starting_points(mask.shape):
161194
if mask[ym, xm] == 0:
162195
xg, yg = dmap.mask2grid(xm, ym)
163-
t = integrate(xg, yg, broken_streamlines)
196+
t = integrate(xg, yg, broken_streamlines,
197+
integration_max_step_scale,
198+
integration_max_error_scale)
164199
if t is not None:
165200
trajectories.append(t)
166201
else:
@@ -188,7 +223,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
188223
xg = np.clip(xg, 0, grid.nx - 1)
189224
yg = np.clip(yg, 0, grid.ny - 1)
190225

191-
t = integrate(xg, yg, broken_streamlines)
226+
t = integrate(xg, yg, broken_streamlines, integration_max_step_scale,
227+
integration_max_error_scale)
192228
if t is not None:
193229
trajectories.append(t)
194230

@@ -481,7 +517,8 @@ def backward_time(xi, yi):
481517
dxi, dyi = forward_time(xi, yi)
482518
return -dxi, -dyi
483519

484-
def integrate(x0, y0, broken_streamlines=True):
520+
def integrate(x0, y0, broken_streamlines=True, integration_max_step_scale=1.0,
521+
integration_max_error_scale=1.0):
485522
"""
486523
Return x, y grid-coordinates of trajectory based on starting point.
487524
@@ -501,14 +538,18 @@ def integrate(x0, y0, broken_streamlines=True):
501538
return None
502539
if integration_direction in ['both', 'backward']:
503540
s, xyt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength,
504-
broken_streamlines)
541+
broken_streamlines,
542+
integration_max_step_scale,
543+
integration_max_error_scale)
505544
stotal += s
506545
xy_traj += xyt[::-1]
507546

508547
if integration_direction in ['both', 'forward']:
509548
dmap.reset_start_point(x0, y0)
510549
s, xyt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength,
511-
broken_streamlines)
550+
broken_streamlines,
551+
integration_max_step_scale,
552+
integration_max_error_scale)
512553
stotal += s
513554
xy_traj += xyt[1:]
514555

@@ -525,7 +566,9 @@ class OutOfBounds(IndexError):
525566
pass
526567

527568

528-
def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True):
569+
def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True,
570+
integration_max_step_scale=1.0,
571+
integration_max_error_scale=1.0):
529572
"""
530573
2nd-order Runge-Kutta algorithm with adaptive step size.
531574
@@ -551,7 +594,7 @@ def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True):
551594
# This error is below that needed to match the RK4 integrator. It
552595
# is set for visual reasons -- too low and corners start
553596
# appearing ugly and jagged. Can be tuned.
554-
maxerror = 0.003
597+
maxerror = 0.003 * integration_max_error_scale
555598

556599
# This limit is important (for all integrators) to avoid the
557600
# trajectory skipping some mask cells. We could relax this
@@ -560,6 +603,7 @@ def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True):
560603
# nature of the interpolation, this doesn't boost speed by much
561604
# for quite a bit of complexity.
562605
maxds = min(1. / dmap.mask.nx, 1. / dmap.mask.ny, 0.1)
606+
maxds *= integration_max_step_scale
563607

564608
ds = maxds
565609
stotal = 0

‎lib/matplotlib/streamplot.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ def streamplot(
2929
integration_direction: Literal["forward", "backward", "both"] = ...,
3030
broken_streamlines: bool = ...,
3131
*,
32+
integration_max_step_scale: float = ...,
33+
integration_max_error_scale: float = ...,
3234
num_arrows: int = ...,
3335
) -> StreamplotSet: ...
3436

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

450 KB
Loading

‎lib/matplotlib/tests/test_streamplot.py

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,50 @@ def test_direction():
100100
linewidth=2, density=2)
101101

102102

103+
@image_comparison(['streamplot_integration.png'], style='mpl20')
104+
def test_integration_options():
105+
# Linear potential flow over a lifting cylinder
106+
n = 50
107+
x, y = np.meshgrid(np.linspace(-2, 2, n), np.linspace(-3, 3, n))
108+
th = np.arctan2(y, x)
109+
r = np.sqrt(x**2 + y**2)
110+
vr = -np.cos(th) / r**2
111+
vt = -np.sin(th) / r**2 - 1 / r
112+
vx = vr * np.cos(th) - vt * np.sin(th) + 1.0
113+
vy = vr * np.sin(th) + vt * np.cos(th)
114+
115+
# Seed points
116+
n_seed = 50
117+
seed_pts = np.column_stack((np.full(n_seed, -1.75), np.linspace(-2, 2, n_seed)))
118+
119+
fig, axs = plt.subplots(3, 1, figsize=(6, 14))
120+
for ax, max_val in zip(axs, [0.05, 1, 5]):
121+
ax.streamplot(
122+
x,
123+
y,
124+
vx,
125+
vy,
126+
start_points=seed_pts,
127+
broken_streamlines=False,
128+
arrowsize=1e-10,
129+
linewidth=0.6,
130+
color="k",
131+
integration_max_step_scale=max_val,
132+
integration_max_error_scale=max_val,
133+
)
134+
135+
# Draw the cylinder
136+
th_circ = np.linspace(0, 2 * np.pi, 100)
137+
for ax in axs:
138+
ax.fill(np.cos(th_circ), np.sin(th_circ), color="w", ec="k")
139+
140+
ax.set_aspect("equal")
141+
ax.set_ylim(-1.5, 1.5)
142+
ax.axis("off")
143+
144+
fig.tight_layout()
145+
146+
103147
def test_streamplot_limits():
104148
ax = plt.axes()
105149
x = np.linspace(-5, 10, 20)
@@ -156,8 +200,20 @@ def test_streamplot_grid():
156200
x = np.array([0, 20, 40])
157201
y = np.array([0, 20, 10])
158202

159-
with pytest.raises(ValueError, match="'y' must be strictly increasing"):
160-
plt.streamplot(x, y, u, v)
203+
204+
def test_streamplot_integration_params():
205+
x = np.array([[10, 20], [10, 20]])
206+
y = np.array([[10, 10], [20, 20]])
207+
u = np.ones((2, 2))
208+
v = np.zeros((2, 2))
209+
210+
err_str = "The value of integration_max_step_scale must be > 0, got -0.5"
211+
with pytest.raises(ValueError, match=err_str):
212+
plt.streamplot(x, y, u, v, integration_max_step_scale=-0.5)
213+
214+
err_str = "The value of integration_max_error_scale must be > 0, got 0.0"
215+
with pytest.raises(ValueError, match=err_str):
216+
plt.streamplot(x, y, u, v, integration_max_error_scale=0.0)
161217

162218

163219
def test_streamplot_inputs(): # test no exception occurs.

0 commit comments

Comments
 (0)