8000 Merge pull request #6062 from gzahl/master · matplotlib/matplotlib@33a1cad · GitHub
[go: up one dir, main page]
Skip to content

Commit 33a1cad

Browse files
NelleVNelleV
authored
Merge pull request #6062 from gzahl/master
Add maximum streamline length property.
2 parents 258dbf5 + 4652f7f commit 33a1cad

File tree

7 files changed

+90
-28
lines changed

7 files changed

10000
+90
-28
lines changed

‎doc/users/whats_new/streamplot_set_maximum_length.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Maximum streamline length and integration direction can now be specified
2+
------------------------------------------------------------------------
3+
4+
This allows to follow the vector field for a longer time and can enhance the
5+
visibility of the flow pattern in some use cases.

‎lib/matplotlib/axes/_axes.py

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4606,21 +4606,25 @@ def stackplot(self, x, *args, **kwargs):
46064606
def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
46074607
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
46084608
minlength=0.1, transform=None, zorder=None,
4609-
start_points=None):
4609+
start_points=None, maxlength=4.0,
4610+
integration_direction='both'):
46104611
if not self._hold:
46114612
self.cla()
4612-
stream_container = mstream.streamplot(self, x, y, u, v,
4613-
density=density,
4614-
linewidth=linewidth,
4615-
color=color,
4616-
cmap=cmap,
4617-
norm=norm,
4618-
arrowsize=arrowsize,
4619-
arrowstyle=arrowstyle,
4620-
minlength=minlength,
4621-
start_points=start_points,
4622-
transform=transform,
4623-
zorder=zorder)
4613+
stream_container = mstream.streamplot(
4614+
self, x, y, u, v,
4615+
density=density,
4616+
linewidth=linewidth,
4617+
color=color,
4618+
cmap=cmap,
4619+
norm=norm,
4620+
arrowsize=arrowsize,
4621+
arrowstyle=arrowstyle,
4622+
minlength=minlength,
4623+
start_points=start_points,
4624+
transform=transform,
4625+
zorder=zorder,
4626+
maxlength=maxlength,
4627+
integration_direction=integration_direction)
46244628
return stream_container
46254629
streamplot.__doc__ = mstream.streamplot.__doc__
46264630

‎lib/matplotlib/pyplot.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3288,8 +3288,8 @@ def step(x, y, *args, **kwargs):
32883288
@_autogen_docstring(Axes.streamplot)
32893289
def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
32903290
norm=None, arrowsize=1, arrowstyle='-|>', minlength=0.1,
3291-
transform=None, zorder=None, start_points=None, hold=None,
3292-
data=None):
3291+
transform=None, zorder=None, start_points=None, maxlength=4.0,
3292+
integration_direction='both', hold=None, data=None):
32933293
ax = gca()
32943294
# allow callers to override the hold state by passing hold=True|False
32953295
washold = ax.ishold()
@@ -3301,7 +3301,10 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
33013301
color=color, cmap=cmap, norm=norm,
33023302
arrowsize=arrowsize, arrowstyle=arrowstyle,
33033303
minlength=minlength, transform=transform,
3304-
zorder=zorder, start_points=start_points, data=data)
3304+
zorder=zorder, start_points=start_points,
3305+
maxlength=maxlength,
3306+
integration_direction=integration_direction,
3307+
data=data)
33053308
finally:
33063309
ax.hold(washold)
33073310
sci(ret.lines)

‎lib/matplotlib/streamplot.py

Lines changed: 37 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@
2222

2323
def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
2424
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
25-
minlength=0.1, transform=None, zorder=None, start_points=None):
25+
minlength=0.1, transform=None, zorder=None, start_points=None,
26+
maxlength=4.0, integration_direction='both'):
2627
"""Draws streamlines of a vector flow.
2728
2829
*x*, *y* : 1d arrays
@@ -58,6 +59,10 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
5859
In data coordinates, the same as the ``x`` and ``y`` arrays.
5960
*zorder* : int
6061
any number
62+
*maxlength* : float
63+
Maximum length of streamline in axes coordinates.
64+
*integration_direction* : ['forward', 'backward', 'both']
65+
Integrate the streamline in forward, backward or both directions.
6166
6267
Returns:
6368
@@ -95,6 +100,15 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
95100
line_kw = {}
96101
arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10 * arrowsize)
97102

103+
if integration_direction not in ['both', 'forward', 'backward']:
104+
errstr = ("Integration direction '%s' not recognised. "
105+
"Expected 'both', 'forward' or 'backward'." %
106+
integration_direction)
107+
raise ValueError(errstr)
108+
109+
if integration_direction == 'both':
110+
maxlength /= 2.
111+
98112
use_multicolor_lines = isinstance(color, np.ndarray)
99113
if use_multicolor_lines:
100114
if color.shape != grid.shape:
@@ -126,7 +140,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
126140
u = np.ma.masked_invalid(u)
127141
v = np.ma.masked_invalid(v)
128142

129-
integrate = get_integrator(u, v, dmap, minlength)
143+
integrate = get_integrator(u, v, dmap, minlength, maxlength,
144+
integration_direction)
130145

131146
trajectories = []
132147
if start_points is None:
@@ -401,7 +416,7 @@ class TerminateTrajectory(Exception):
401416
# Integrator definitions
402417
#========================
403418

404-
def get_integrator(u, v, dmap, minlength):
419+
def get_integrator(u, v, dmap, minlength, maxlength, integration_direction):
405420

406421
# rescale velocity onto grid-coordinates for integrations.
407422
u, v = dmap.data2grid(u, v)
@@ -435,17 +450,27 @@ def integrate(x0, y0):
435450
resulting trajectory is None if it is shorter than `minlength`.
436451
"""
437452

453+
stotal, x_traj, y_traj = 0., [], []
454+
438455
try:
439456
dmap.start_trajectory(x0, y0)
440457
except InvalidIndexError:
441458
return None
442-
sf, xf_traj, yf_traj = _integrate_rk12(x0, y0, dmap, forward_time)
443-
dmap.reset_start_point(x0, y0)
444-
sb, xb_traj, yb_traj = _integrate_rk12(x0, y0, dmap, backward_time)
445-
# combine forward and backward trajectories
446-
stotal = sf + sb
447-
x_traj = xb_traj[::-1] + xf_traj[1:]
448-
y_traj = yb_traj[::-1] + yf_traj[1:]
459+
if integration_direction in ['both', 'backward']:
460+
s, xt, yt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength)
461+
stotal += s
462+
x_traj += xt[::-1]
463+
y_traj += yt[::-1]
464+
465+
if integration_direction in ['both', 'forward']:
466+
dmap.reset_start_point(x0, y0)
467+
s, xt, yt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength)
468+
if len(x_traj) > 0:
469+
xt = xt[1:]
470+
yt = yt[1:]
471+
stotal += s
472+
x_traj += xt
473+
y_traj += yt
449474

450475
if stotal > minlength:
451476
return x_traj, y_traj
@@ -456,7 +481,7 @@ def integrate(x0, y0):
456481
return integrate
457482

458483

459-
def _integrate_rk12(x0, y0, dmap, f):
484+
def _integrate_rk12(x0, y0, dmap, f, maxlength):
460485
"""2nd-order Runge-Kutta algorithm with adaptive step size.
461486
462487
This method is also referred to as the improved Euler's method, or Heun's
@@ -532,7 +557,7 @@ def _integrate_rk12(x0, y0, dmap, f):
532557
dmap.update_trajectory(xi, yi)
533558
except InvalidIndexError:
534559
break
535-
if (stotal + ds) > 2:
560+
if (stotal + ds) > maxlength:
536561
break
537562
stotal += ds
538563

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

26.2 KB
Loading

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

39 KB
Loading

‎lib/matplotlib/tests/test_streamplot.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,14 @@ def velocity_field():
1616
V = 1 + X - Y**2
1717
return X, Y, U, V
1818

19+
def swirl_velocity_field():
20+
x = np.linspace(-3., 3., 100)
21+
y = np.linspace(-3., 3., 100)
22+
X, Y = np.meshgrid(x, y)
23+
a = 0.1
24+
U = np.cos(a) * (-Y) - np.sin(a) * X
25+
V = np.sin(a) * (-Y) + np.cos(a) * X
26+
return x, y, U, V
1927

2028
@image_comparison(baseline_images=['streamplot_startpoints'])
2129
def test_startpoints():
@@ -57,6 +65,23 @@ def test_masks_and_nans():
5765
plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)
5866

5967

68+
@image_comparison(baseline_images=['streamplot_maxlength'],
69+
extensions=['png'])
70+
def test_maxlength():
71+
x, y, U, V = swirl_velocity_field()
72+
plt.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
73+
linewidth=2, density=2)
74+
75+
76+
@image_comparison(baseline_images=['streamplot_direction'],
77+
extensions=['png'])
78+
def test_direction():
79+
x, y, U, V = swirl_velocity_field()
80+
plt.streamplot(x, y, U, V, integration_direction='backward',
81+
maxlength=1.5, start_points=[[1.5, 0.]],
82+
linewidth=2, density=2)
83+
84+
6085
@cleanup
6186
def test_streamplot_limits():
6287
ax = plt.axes()

0 commit comments

Comments
 (0)
0