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

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 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

+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
A851 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