Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/matplotlib/streamplot.py: 11%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Streamline plotting for 2D vector fields.
4"""
6import numpy as np
8import matplotlib as mpl
9from matplotlib import _api, cm, patches
10import matplotlib.colors as mcolors
11import matplotlib.collections as mcollections
12import matplotlib.lines as mlines
15__all__ = ['streamplot']
18def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
19 cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
20 minlength=0.1, transform=None, zorder=None, start_points=None,
21 maxlength=4.0, integration_direction='both',
22 broken_streamlines=True):
23 """
24 Draw streamlines of a vector flow.
26 Parameters
27 ----------
28 x, y : 1D/2D arrays
29 Evenly spaced strictly increasing arrays to make a grid. If 2D, all
30 rows of *x* must be equal and all columns of *y* must be equal; i.e.,
31 they must be as if generated by ``np.meshgrid(x_1d, y_1d)``.
32 u, v : 2D arrays
33 *x* and *y*-velocities. The number of rows and columns must match
34 the length of *y* and *x*, respectively.
35 density : float or (float, float)
36 Controls the closeness of streamlines. When ``density = 1``, the domain
37 is divided into a 30x30 grid. *density* linearly scales this grid.
38 Each cell in the grid can have, at most, one traversing streamline.
39 For different densities in each direction, use a tuple
40 (density_x, density_y).
41 linewidth : float or 2D array
42 The width of the streamlines. With a 2D array the line width can be
43 varied across the grid. The array must have the same shape as *u*
44 and *v*.
45 color : :mpltype:`color` or 2D array
46 The streamline color. If given an array, its values are converted to
47 colors using *cmap* and *norm*. The array must have the same shape
48 as *u* and *v*.
49 cmap, norm
50 Data normalization and colormapping parameters for *color*; only used
51 if *color* is an array of floats. See `~.Axes.imshow` for a detailed
52 description.
53 arrowsize : float
54 Scaling factor for the arrow size.
55 arrowstyle : str
56 Arrow style specification.
57 See `~matplotlib.patches.FancyArrowPatch`.
58 minlength : float
59 Minimum length of streamline in axes coordinates.
60 start_points : (N, 2) array
61 Coordinates of starting points for the streamlines in data coordinates
62 (the same coordinates as the *x* and *y* arrays).
63 zorder : float
64 The zorder of the streamlines and arrows.
65 Artists with lower zorder values are drawn first.
66 maxlength : float
67 Maximum length of streamline in axes coordinates.
68 integration_direction : {'forward', 'backward', 'both'}, default: 'both'
69 Integrate the streamline in forward, backward or both directions.
70 data : indexable object, optional
71 DATA_PARAMETER_PLACEHOLDER
72 broken_streamlines : boolean, default: True
73 If False, forces streamlines to continue until they
74 leave the plot domain. If True, they may be terminated if they
75 come too close to another streamline.
77 Returns
78 -------
79 StreamplotSet
80 Container object with attributes
82 - ``lines``: `.LineCollection` of streamlines
84 - ``arrows``: `.PatchCollection` containing `.FancyArrowPatch`
85 objects representing the arrows half-way along streamlines.
87 This container will probably change in the future to allow changes
88 to the colormap, alpha, etc. for both lines and arrows, but these
89 changes should be backward compatible.
90 """
91 grid = Grid(x, y)
92 mask = StreamMask(density)
93 dmap = DomainMap(grid, mask)
95 if zorder is None:
96 zorder = mlines.Line2D.zorder
98 # default to data coordinates
99 if transform is None:
100 transform = axes.transData
102 if color is None:
103 color = axes._get_lines.get_next_color()
105 if linewidth is None:
106 linewidth = mpl.rcParams['lines.linewidth']
108 line_kw = {}
109 arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10 * arrowsize)
111 _api.check_in_list(['both', 'forward', 'backward'],
112 integration_direction=integration_direction)
114 if integration_direction == 'both':
115 maxlength /= 2.
117 use_multicolor_lines = isinstance(color, np.ndarray)
118 if use_multicolor_lines:
119 if color.shape != grid.shape:
120 raise ValueError("If 'color' is given, it must match the shape of "
121 "the (x, y) grid")
122 line_colors = [[]] # Empty entry allows concatenation of zero arrays.
123 color = np.ma.masked_invalid(color)
124 else:
125 line_kw['color'] = color
126 arrow_kw['color'] = color
128 if isinstance(linewidth, np.ndarray):
129 if linewidth.shape != grid.shape:
130 raise ValueError("If 'linewidth' is given, it must match the "
131 "shape of the (x, y) grid")
132 line_kw['linewidth'] = []
133 else:
134 line_kw['linewidth'] = linewidth
135 arrow_kw['linewidth'] = linewidth
137 line_kw['zorder'] = zorder
138 arrow_kw['zorder'] = zorder
140 # Sanity checks.
141 if u.shape != grid.shape or v.shape != grid.shape:
142 raise ValueError("'u' and 'v' must match the shape of the (x, y) grid")
144 u = np.ma.masked_invalid(u)
145 v = np.ma.masked_invalid(v)
147 integrate = _get_integrator(u, v, dmap, minlength, maxlength,
148 integration_direction)
150 trajectories = []
151 if start_points is None:
152 for xm, ym in _gen_starting_points(mask.shape):
153 if mask[ym, xm] == 0:
154 xg, yg = dmap.mask2grid(xm, ym)
155 t = integrate(xg, yg, broken_streamlines)
156 if t is not None:
157 trajectories.append(t)
158 else:
159 sp2 = np.asanyarray(start_points, dtype=float).copy()
161 # Check if start_points are outside the data boundaries
162 for xs, ys in sp2:
163 if not (grid.x_origin <= xs <= grid.x_origin + grid.width and
164 grid.y_origin <= ys <= grid.y_origin + grid.height):
165 raise ValueError(f"Starting point ({xs}, {ys}) outside of "
166 "data boundaries")
168 # Convert start_points from data to array coords
169 # Shift the seed points from the bottom left of the data so that
170 # data2grid works properly.
171 sp2[:, 0] -= grid.x_origin
172 sp2[:, 1] -= grid.y_origin
174 for xs, ys in sp2:
175 xg, yg = dmap.data2grid(xs, ys)
176 # Floating point issues can cause xg, yg to be slightly out of
177 # bounds for xs, ys on the upper boundaries. Because we have
178 # already checked that the starting points are within the original
179 # grid, clip the xg, yg to the grid to work around this issue
180 xg = np.clip(xg, 0, grid.nx - 1)
181 yg = np.clip(yg, 0, grid.ny - 1)
183 t = integrate(xg, yg, broken_streamlines)
184 if t is not None:
185 trajectories.append(t)
187 if use_multicolor_lines:
188 if norm is None:
189 norm = mcolors.Normalize(color.min(), color.max())
190 cmap = cm._ensure_cmap(cmap)
192 streamlines = []
193 arrows = []
194 for t in trajectories:
195 tgx, tgy = t.T
196 # Rescale from grid-coordinates to data-coordinates.
197 tx, ty = dmap.grid2data(tgx, tgy)
198 tx += grid.x_origin
199 ty += grid.y_origin
201 # Create multiple tiny segments if varying width or color is given
202 if isinstance(linewidth, np.ndarray) or use_multicolor_lines:
203 points = np.transpose([tx, ty]).reshape(-1, 1, 2)
204 streamlines.extend(np.hstack([points[:-1], points[1:]]))
205 else:
206 points = np.transpose([tx, ty])
207 streamlines.append(points)
209 # Add arrows halfway along each trajectory.
210 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]))
215 if isinstance(linewidth, np.ndarray):
216 line_widths = interpgrid(linewidth, tgx, tgy)[:-1]
217 line_kw['linewidth'].extend(line_widths)
218 arrow_kw['linewidth'] = line_widths[n]
220 if use_multicolor_lines:
221 color_values = interpgrid(color, tgx, tgy)[:-1]
222 line_colors.append(color_values)
223 arrow_kw['color'] = cmap(norm(color_values[n]))
225 p = patches.FancyArrowPatch(
226 arrow_tail, arrow_head, transform=transform, **arrow_kw)
227 arrows.append(p)
229 lc = mcollections.LineCollection(
230 streamlines, transform=transform, **line_kw)
231 lc.sticky_edges.x[:] = [grid.x_origin, grid.x_origin + grid.width]
232 lc.sticky_edges.y[:] = [grid.y_origin, grid.y_origin + grid.height]
233 if use_multicolor_lines:
234 lc.set_array(np.ma.hstack(line_colors))
235 lc.set_cmap(cmap)
236 lc.set_norm(norm)
237 axes.add_collection(lc)
239 ac = mcollections.PatchCollection(arrows)
240 # Adding the collection itself is broken; see #2341.
241 for p in arrows:
242 axes.add_patch(p)
244 axes.autoscale_view()
245 stream_container = StreamplotSet(lc, ac)
246 return stream_container
249class StreamplotSet:
251 def __init__(self, lines, arrows):
252 self.lines = lines
253 self.arrows = arrows
256# Coordinate definitions
257# ========================
259class DomainMap:
260 """
261 Map representing different coordinate systems.
263 Coordinate definitions:
265 * axes-coordinates goes from 0 to 1 in the domain.
266 * data-coordinates are specified by the input x-y coordinates.
267 * grid-coordinates goes from 0 to N and 0 to M for an N x M grid,
268 where N and M match the shape of the input data.
269 * mask-coordinates goes from 0 to N and 0 to M for an N x M mask,
270 where N and M are user-specified to control the density of streamlines.
272 This class also has methods for adding trajectories to the StreamMask.
273 Before adding a trajectory, run `start_trajectory` to keep track of regions
274 crossed by a given trajectory. Later, if you decide the trajectory is bad
275 (e.g., if the trajectory is very short) just call `undo_trajectory`.
276 """
278 def __init__(self, grid, mask):
279 self.grid = grid
280 self.mask = mask
281 # Constants for conversion between grid- and mask-coordinates
282 self.x_grid2mask = (mask.nx - 1) / (grid.nx - 1)
283 self.y_grid2mask = (mask.ny - 1) / (grid.ny - 1)
285 self.x_mask2grid = 1. / self.x_grid2mask
286 self.y_mask2grid = 1. / self.y_grid2mask
288 self.x_data2grid = 1. / grid.dx
289 self.y_data2grid = 1. / grid.dy
291 def grid2mask(self, xi, yi):
292 """Return nearest space in mask-coords from given grid-coords."""
293 return round(xi * self.x_grid2mask), round(yi * self.y_grid2mask)
295 def mask2grid(self, xm, ym):
296 return xm * self.x_mask2grid, ym * self.y_mask2grid
298 def data2grid(self, xd, yd):
299 return xd * self.x_data2grid, yd * self.y_data2grid
301 def grid2data(self, xg, yg):
302 return xg / self.x_data2grid, yg / self.y_data2grid
304 def start_trajectory(self, xg, yg, broken_streamlines=True):
305 xm, ym = self.grid2mask(xg, yg)
306 self.mask._start_trajectory(xm, ym, broken_streamlines)
308 def reset_start_point(self, xg, yg):
309 xm, ym = self.grid2mask(xg, yg)
310 self.mask._current_xy = (xm, ym)
312 def update_trajectory(self, xg, yg, broken_streamlines=True):
313 if not self.grid.within_grid(xg, yg):
314 raise InvalidIndexError
315 xm, ym = self.grid2mask(xg, yg)
316 self.mask._update_trajectory(xm, ym, broken_streamlines)
318 def undo_trajectory(self):
319 self.mask._undo_trajectory()
322class Grid:
323 """Grid of data."""
324 def __init__(self, x, y):
326 if np.ndim(x) == 1:
327 pass
328 elif np.ndim(x) == 2:
329 x_row = x[0]
330 if not np.allclose(x_row, x):
331 raise ValueError("The rows of 'x' must be equal")
332 x = x_row
333 else:
334 raise ValueError("'x' can have at maximum 2 dimensions")
336 if np.ndim(y) == 1:
337 pass
338 elif np.ndim(y) == 2:
339 yt = np.transpose(y) # Also works for nested lists.
340 y_col = yt[0]
341 if not np.allclose(y_col, yt):
342 raise ValueError("The columns of 'y' must be equal")
343 y = y_col
344 else:
345 raise ValueError("'y' can have at maximum 2 dimensions")
347 if not (np.diff(x) > 0).all():
348 raise ValueError("'x' must be strictly increasing")
349 if not (np.diff(y) > 0).all():
350 raise ValueError("'y' must be strictly increasing")
352 self.nx = len(x)
353 self.ny = len(y)
355 self.dx = x[1] - x[0]
356 self.dy = y[1] - y[0]
358 self.x_origin = x[0]
359 self.y_origin = y[0]
361 self.width = x[-1] - x[0]
362 self.height = y[-1] - y[0]
364 if not np.allclose(np.diff(x), self.width / (self.nx - 1)):
365 raise ValueError("'x' values must be equally spaced")
366 if not np.allclose(np.diff(y), self.height / (self.ny - 1)):
367 raise ValueError("'y' values must be equally spaced")
369 @property
370 def shape(self):
371 return self.ny, self.nx
373 def within_grid(self, xi, yi):
374 """Return whether (*xi*, *yi*) is a valid index of the grid."""
375 # Note that xi/yi can be floats; so, for example, we can't simply check
376 # `xi < self.nx` since *xi* can be `self.nx - 1 < xi < self.nx`
377 return 0 <= xi <= self.nx - 1 and 0 <= yi <= self.ny - 1
380class StreamMask:
381 """
382 Mask to keep track of discrete regions crossed by streamlines.
384 The resolution of this grid determines the approximate spacing between
385 trajectories. Streamlines are only allowed to pass through zeroed cells:
386 When a streamline enters a cell, that cell is set to 1, and no new
387 streamlines are allowed to enter.
388 """
390 def __init__(self, density):
391 try:
392 self.nx, self.ny = (30 * np.broadcast_to(density, 2)).astype(int)
393 except ValueError as err:
394 raise ValueError("'density' must be a scalar or be of length "
395 "2") from err
396 if self.nx < 0 or self.ny < 0:
397 raise ValueError("'density' must be positive")
398 self._mask = np.zeros((self.ny, self.nx))
399 self.shape = self._mask.shape
401 self._current_xy = None
403 def __getitem__(self, args):
404 return self._mask[args]
406 def _start_trajectory(self, xm, ym, broken_streamlines=True):
407 """Start recording streamline trajectory"""
408 self._traj = []
409 self._update_trajectory(xm, ym, broken_streamlines)
411 def _undo_trajectory(self):
412 """Remove current trajectory from mask"""
413 for t in self._traj:
414 self._mask[t] = 0
416 def _update_trajectory(self, xm, ym, broken_streamlines=True):
417 """
418 Update current trajectory position in mask.
420 If the new position has already been filled, raise `InvalidIndexError`.
421 """
422 if self._current_xy != (xm, ym):
423 if self[ym, xm] == 0:
424 self._traj.append((ym, xm))
425 self._mask[ym, xm] = 1
426 self._current_xy = (xm, ym)
427 else:
428 if broken_streamlines:
429 raise InvalidIndexError
430 else:
431 pass
434class InvalidIndexError(Exception):
435 pass
438class TerminateTrajectory(Exception):
439 pass
442# Integrator definitions
443# =======================
445def _get_integrator(u, v, dmap, minlength, maxlength, integration_direction):
447 # rescale velocity onto grid-coordinates for integrations.
448 u, v = dmap.data2grid(u, v)
450 # speed (path length) will be in axes-coordinates
451 u_ax = u / (dmap.grid.nx - 1)
452 v_ax = v / (dmap.grid.ny - 1)
453 speed = np.ma.sqrt(u_ax ** 2 + v_ax ** 2)
455 def forward_time(xi, yi):
456 if not dmap.grid.within_grid(xi, yi):
457 raise OutOfBounds
458 ds_dt = interpgrid(speed, xi, yi)
459 if ds_dt == 0:
460 raise TerminateTrajectory()
461 dt_ds = 1. / ds_dt
462 ui = interpgrid(u, xi, yi)
463 vi = interpgrid(v, xi, yi)
464 return ui * dt_ds, vi * dt_ds
466 def backward_time(xi, yi):
467 dxi, dyi = forward_time(xi, yi)
468 return -dxi, -dyi
470 def integrate(x0, y0, broken_streamlines=True):
471 """
472 Return x, y grid-coordinates of trajectory based on starting point.
474 Integrate both forward and backward in time from starting point in
475 grid coordinates.
477 Integration is terminated when a trajectory reaches a domain boundary
478 or when it crosses into an already occupied cell in the StreamMask. The
479 resulting trajectory is None if it is shorter than `minlength`.
480 """
482 stotal, xy_traj = 0., []
484 try:
485 dmap.start_trajectory(x0, y0, broken_streamlines)
486 except InvalidIndexError:
487 return None
488 if integration_direction in ['both', 'backward']:
489 s, xyt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength,
490 broken_streamlines)
491 stotal += s
492 xy_traj += xyt[::-1]
494 if integration_direction in ['both', 'forward']:
495 dmap.reset_start_point(x0, y0)
496 s, xyt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength,
497 broken_streamlines)
498 stotal += s
499 xy_traj += xyt[1:]
501 if stotal > minlength:
502 return np.broadcast_arrays(xy_traj, np.empty((1, 2)))[0]
503 else: # reject short trajectories
504 dmap.undo_trajectory()
505 return None
507 return integrate
510class OutOfBounds(IndexError):
511 pass
514def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True):
515 """
516 2nd-order Runge-Kutta algorithm with adaptive step size.
518 This method is also referred to as the improved Euler's method, or Heun's
519 method. This method is favored over higher-order methods because:
521 1. To get decent looking trajectories and to sample every mask cell
522 on the trajectory we need a small timestep, so a lower order
523 solver doesn't hurt us unless the data is *very* high resolution.
524 In fact, for cases where the user inputs
525 data smaller or of similar grid size to the mask grid, the higher
526 order corrections are negligible because of the very fast linear
527 interpolation used in `interpgrid`.
529 2. For high resolution input data (i.e. beyond the mask
530 resolution), we must reduce the timestep. Therefore, an adaptive
531 timestep is more suited to the problem as this would be very hard
532 to judge automatically otherwise.
534 This integrator is about 1.5 - 2x as fast as RK4 and RK45 solvers (using
535 similar Python implementations) in most setups.
536 """
537 # This error is below that needed to match the RK4 integrator. It
538 # is set for visual reasons -- too low and corners start
539 # appearing ugly and jagged. Can be tuned.
540 maxerror = 0.003
542 # This limit is important (for all integrators) to avoid the
543 # trajectory skipping some mask cells. We could relax this
544 # condition if we use the code which is commented out below to
545 # increment the location gradually. However, due to the efficient
546 # nature of the interpolation, this doesn't boost speed by much
547 # for quite a bit of complexity.
548 maxds = min(1. / dmap.mask.nx, 1. / dmap.mask.ny, 0.1)
550 ds = maxds
551 stotal = 0
552 xi = x0
553 yi = y0
554 xyf_traj = []
556 while True:
557 try:
558 if dmap.grid.within_grid(xi, yi):
559 xyf_traj.append((xi, yi))
560 else:
561 raise OutOfBounds
563 # Compute the two intermediate gradients.
564 # f should raise OutOfBounds if the locations given are
565 # outside the grid.
566 k1x, k1y = f(xi, yi)
567 k2x, k2y = f(xi + ds * k1x, yi + ds * k1y)
569 except OutOfBounds:
570 # Out of the domain during this step.
571 # Take an Euler step to the boundary to improve neatness
572 # unless the trajectory is currently empty.
573 if xyf_traj:
574 ds, xyf_traj = _euler_step(xyf_traj, dmap, f)
575 stotal += ds
576 break
577 except TerminateTrajectory:
578 break
580 dx1 = ds * k1x
581 dy1 = ds * k1y
582 dx2 = ds * 0.5 * (k1x + k2x)
583 dy2 = ds * 0.5 * (k1y + k2y)
585 ny, nx = dmap.grid.shape
586 # Error is normalized to the axes coordinates
587 error = np.hypot((dx2 - dx1) / (nx - 1), (dy2 - dy1) / (ny - 1))
589 # Only save step if within error tolerance
590 if error < maxerror:
591 xi += dx2
592 yi += dy2
593 try:
594 dmap.update_trajectory(xi, yi, broken_streamlines)
595 except InvalidIndexError:
596 break
597 if stotal + ds > maxlength:
598 break
599 stotal += ds
601 # recalculate stepsize based on step error
602 if error == 0:
603 ds = maxds
604 else:
605 ds = min(maxds, 0.85 * ds * (maxerror / error) ** 0.5)
607 return stotal, xyf_traj
610def _euler_step(xyf_traj, dmap, f):
611 """Simple Euler integration step that extends streamline to boundary."""
612 ny, nx = dmap.grid.shape
613 xi, yi = xyf_traj[-1]
614 cx, cy = f(xi, yi)
615 if cx == 0:
616 dsx = np.inf
617 elif cx < 0:
618 dsx = xi / -cx
619 else:
620 dsx = (nx - 1 - xi) / cx
621 if cy == 0:
622 dsy = np.inf
623 elif cy < 0:
624 dsy = yi / -cy
625 else:
626 dsy = (ny - 1 - yi) / cy
627 ds = min(dsx, dsy)
628 xyf_traj.append((xi + cx * ds, yi + cy * ds))
629 return ds, xyf_traj
632# Utility functions
633# ========================
635def interpgrid(a, xi, yi):
636 """Fast 2D, linear interpolation on an integer grid"""
638 Ny, Nx = np.shape(a)
639 if isinstance(xi, np.ndarray):
640 x = xi.astype(int)
641 y = yi.astype(int)
642 # Check that xn, yn don't exceed max index
643 xn = np.clip(x + 1, 0, Nx - 1)
644 yn = np.clip(y + 1, 0, Ny - 1)
645 else:
646 x = int(xi)
647 y = int(yi)
648 # conditional is faster than clipping for integers
649 if x == (Nx - 1):
650 xn = x
651 else:
652 xn = x + 1
653 if y == (Ny - 1):
654 yn = y
655 else:
656 yn = y + 1
658 a00 = a[y, x]
659 a01 = a[y, xn]
660 a10 = a[yn, x]
661 a11 = a[yn, xn]
662 xt = xi - x
663 yt = yi - y
664 a0 = a00 * (1 - xt) + a01 * xt
665 a1 = a10 * (1 - xt) + a11 * xt
666 ai = a0 * (1 - yt) + a1 * yt
668 if not isinstance(xi, np.ndarray):
669 if np.ma.is_masked(ai):
670 raise TerminateTrajectory
672 return ai
675def _gen_starting_points(shape):
676 """
677 Yield starting points for streamlines.
679 Trying points on the boundary first gives higher quality streamlines.
680 This algorithm starts with a point on the mask corner and spirals inward.
681 This algorithm is inefficient, but fast compared to rest of streamplot.
682 """
683 ny, nx = shape
684 xfirst = 0
685 yfirst = 1
686 xlast = nx - 1
687 ylast = ny - 1
688 x, y = 0, 0
689 direction = 'right'
690 for i in range(nx * ny):
691 yield x, y
693 if direction == 'right':
694 x += 1
695 if x >= xlast:
696 xlast -= 1
697 direction = 'up'
698 elif direction == 'up':
699 y += 1
700 if y >= ylast:
701 ylast -= 1
702 direction = 'left'
703 elif direction == 'left':
704 x -= 1
705 if x <= xfirst:
706 xfirst += 1
707 direction = 'down'
708 elif direction == 'down':
709 y -= 1
710 if y <= yfirst:
711 yfirst += 1
712 direction = 'right'