Fixed database typo and removed unnecessary class identifier.
This commit is contained in:
parent
00ad49a143
commit
45fb349a7d
5098 changed files with 952558 additions and 85 deletions
703
venv/Lib/site-packages/matplotlib/streamplot.py
Normal file
703
venv/Lib/site-packages/matplotlib/streamplot.py
Normal file
|
@ -0,0 +1,703 @@
|
|||
"""
|
||||
Streamline plotting for 2D vector fields.
|
||||
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
import matplotlib
|
||||
import matplotlib.cbook as cbook
|
||||
import matplotlib.cm as cm
|
||||
import matplotlib.colors as mcolors
|
||||
import matplotlib.collections as mcollections
|
||||
import matplotlib.lines as mlines
|
||||
import matplotlib.patches as patches
|
||||
|
||||
|
||||
__all__ = ['streamplot']
|
||||
|
||||
|
||||
def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
|
||||
cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
|
||||
minlength=0.1, transform=None, zorder=None, start_points=None,
|
||||
maxlength=4.0, integration_direction='both'):
|
||||
"""
|
||||
Draw streamlines of a vector flow.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x, y : 1D arrays
|
||||
An evenly spaced grid.
|
||||
u, v : 2D arrays
|
||||
*x* and *y*-velocities. The number of rows and columns must match
|
||||
the length of *y* and *x*, respectively.
|
||||
density : float or (float, float)
|
||||
Controls the closeness of streamlines. When ``density = 1``, the domain
|
||||
is divided into a 30x30 grid. *density* linearly scales this grid.
|
||||
Each cell in the grid can have, at most, one traversing streamline.
|
||||
For different densities in each direction, use a tuple
|
||||
(density_x, density_y).
|
||||
linewidth : float or 2D array
|
||||
The width of the stream lines. With a 2D array the line width can be
|
||||
varied across the grid. The array must have the same shape as *u*
|
||||
and *v*.
|
||||
color : color or 2D array
|
||||
The streamline color. If given an array, its values are converted to
|
||||
colors using *cmap* and *norm*. The array must have the same shape
|
||||
as *u* and *v*.
|
||||
cmap : `~matplotlib.colors.Colormap`
|
||||
Colormap used to plot streamlines and arrows. This is only used if
|
||||
*color* is an array.
|
||||
norm : `~matplotlib.colors.Normalize`
|
||||
Normalize object used to scale luminance data to 0, 1. If ``None``,
|
||||
stretch (min, max) to (0, 1). This is only used if *color* is an array.
|
||||
arrowsize : float
|
||||
Scaling factor for the arrow size.
|
||||
arrowstyle : str
|
||||
Arrow style specification.
|
||||
See `~matplotlib.patches.FancyArrowPatch`.
|
||||
minlength : float
|
||||
Minimum length of streamline in axes coordinates.
|
||||
start_points : Nx2 array
|
||||
Coordinates of starting points for the streamlines in data coordinates
|
||||
(the same coordinates as the *x* and *y* arrays).
|
||||
zorder : int
|
||||
The zorder of the stream lines and arrows.
|
||||
Artists with lower zorder values are drawn first.
|
||||
maxlength : float
|
||||
Maximum length of streamline in axes coordinates.
|
||||
integration_direction : {'forward', 'backward', 'both'}, default: 'both'
|
||||
Integrate the streamline in forward, backward or both directions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
StreamplotSet
|
||||
Container object with attributes
|
||||
|
||||
- ``lines``: `.LineCollection` of streamlines
|
||||
|
||||
- ``arrows``: `.PatchCollection` containing `.FancyArrowPatch`
|
||||
objects representing the arrows half-way along stream lines.
|
||||
|
||||
This container will probably change in the future to allow changes
|
||||
to the colormap, alpha, etc. for both lines and arrows, but these
|
||||
changes should be backward compatible.
|
||||
"""
|
||||
grid = Grid(x, y)
|
||||
mask = StreamMask(density)
|
||||
dmap = DomainMap(grid, mask)
|
||||
|
||||
if zorder is None:
|
||||
zorder = mlines.Line2D.zorder
|
||||
|
||||
# default to data coordinates
|
||||
if transform is None:
|
||||
transform = axes.transData
|
||||
|
||||
if color is None:
|
||||
color = axes._get_lines.get_next_color()
|
||||
|
||||
if linewidth is None:
|
||||
linewidth = matplotlib.rcParams['lines.linewidth']
|
||||
|
||||
line_kw = {}
|
||||
arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10 * arrowsize)
|
||||
|
||||
cbook._check_in_list(['both', 'forward', 'backward'],
|
||||
integration_direction=integration_direction)
|
||||
|
||||
if integration_direction == 'both':
|
||||
maxlength /= 2.
|
||||
|
||||
use_multicolor_lines = isinstance(color, np.ndarray)
|
||||
if use_multicolor_lines:
|
||||
if color.shape != grid.shape:
|
||||
raise ValueError("If 'color' is given, it must match the shape of "
|
||||
"'Grid(x, y)'")
|
||||
line_colors = []
|
||||
color = np.ma.masked_invalid(color)
|
||||
else:
|
||||
line_kw['color'] = color
|
||||
arrow_kw['color'] = color
|
||||
|
||||
if isinstance(linewidth, np.ndarray):
|
||||
if linewidth.shape != grid.shape:
|
||||
raise ValueError("If 'linewidth' is given, it must match the "
|
||||
"shape of 'Grid(x, y)'")
|
||||
line_kw['linewidth'] = []
|
||||
else:
|
||||
line_kw['linewidth'] = linewidth
|
||||
arrow_kw['linewidth'] = linewidth
|
||||
|
||||
line_kw['zorder'] = zorder
|
||||
arrow_kw['zorder'] = zorder
|
||||
|
||||
# Sanity checks.
|
||||
if u.shape != grid.shape or v.shape != grid.shape:
|
||||
raise ValueError("'u' and 'v' must match the shape of 'Grid(x, y)'")
|
||||
|
||||
u = np.ma.masked_invalid(u)
|
||||
v = np.ma.masked_invalid(v)
|
||||
|
||||
integrate = get_integrator(u, v, dmap, minlength, maxlength,
|
||||
integration_direction)
|
||||
|
||||
trajectories = []
|
||||
if start_points is None:
|
||||
for xm, ym in _gen_starting_points(mask.shape):
|
||||
if mask[ym, xm] == 0:
|
||||
xg, yg = dmap.mask2grid(xm, ym)
|
||||
t = integrate(xg, yg)
|
||||
if t is not None:
|
||||
trajectories.append(t)
|
||||
else:
|
||||
sp2 = np.asanyarray(start_points, dtype=float).copy()
|
||||
|
||||
# Check if start_points are outside the data boundaries
|
||||
for xs, ys in sp2:
|
||||
if not (grid.x_origin <= xs <= grid.x_origin + grid.width and
|
||||
grid.y_origin <= ys <= grid.y_origin + grid.height):
|
||||
raise ValueError("Starting point ({}, {}) outside of data "
|
||||
"boundaries".format(xs, ys))
|
||||
|
||||
# Convert start_points from data to array coords
|
||||
# Shift the seed points from the bottom left of the data so that
|
||||
# data2grid works properly.
|
||||
sp2[:, 0] -= grid.x_origin
|
||||
sp2[:, 1] -= grid.y_origin
|
||||
|
||||
for xs, ys in sp2:
|
||||
xg, yg = dmap.data2grid(xs, ys)
|
||||
t = integrate(xg, yg)
|
||||
if t is not None:
|
||||
trajectories.append(t)
|
||||
|
||||
if use_multicolor_lines:
|
||||
if norm is None:
|
||||
norm = mcolors.Normalize(color.min(), color.max())
|
||||
if cmap is None:
|
||||
cmap = cm.get_cmap(matplotlib.rcParams['image.cmap'])
|
||||
else:
|
||||
cmap = cm.get_cmap(cmap)
|
||||
|
||||
streamlines = []
|
||||
arrows = []
|
||||
for t in trajectories:
|
||||
tgx = np.array(t[0])
|
||||
tgy = np.array(t[1])
|
||||
# Rescale from grid-coordinates to data-coordinates.
|
||||
tx, ty = dmap.grid2data(*np.array(t))
|
||||
tx += grid.x_origin
|
||||
ty += grid.y_origin
|
||||
|
||||
points = np.transpose([tx, ty]).reshape(-1, 1, 2)
|
||||
streamlines.extend(np.hstack([points[:-1], points[1:]]))
|
||||
|
||||
# Add arrows half way along each trajectory.
|
||||
s = np.cumsum(np.hypot(np.diff(tx), np.diff(ty)))
|
||||
n = np.searchsorted(s, s[-1] / 2.)
|
||||
arrow_tail = (tx[n], ty[n])
|
||||
arrow_head = (np.mean(tx[n:n + 2]), np.mean(ty[n:n + 2]))
|
||||
|
||||
if isinstance(linewidth, np.ndarray):
|
||||
line_widths = interpgrid(linewidth, tgx, tgy)[:-1]
|
||||
line_kw['linewidth'].extend(line_widths)
|
||||
arrow_kw['linewidth'] = line_widths[n]
|
||||
|
||||
if use_multicolor_lines:
|
||||
color_values = interpgrid(color, tgx, tgy)[:-1]
|
||||
line_colors.append(color_values)
|
||||
arrow_kw['color'] = cmap(norm(color_values[n]))
|
||||
|
||||
p = patches.FancyArrowPatch(
|
||||
arrow_tail, arrow_head, transform=transform, **arrow_kw)
|
||||
axes.add_patch(p)
|
||||
arrows.append(p)
|
||||
|
||||
lc = mcollections.LineCollection(
|
||||
streamlines, transform=transform, **line_kw)
|
||||
lc.sticky_edges.x[:] = [grid.x_origin, grid.x_origin + grid.width]
|
||||
lc.sticky_edges.y[:] = [grid.y_origin, grid.y_origin + grid.height]
|
||||
if use_multicolor_lines:
|
||||
lc.set_array(np.ma.hstack(line_colors))
|
||||
lc.set_cmap(cmap)
|
||||
lc.set_norm(norm)
|
||||
axes.add_collection(lc)
|
||||
axes.autoscale_view()
|
||||
|
||||
ac = matplotlib.collections.PatchCollection(arrows)
|
||||
stream_container = StreamplotSet(lc, ac)
|
||||
return stream_container
|
||||
|
||||
|
||||
class StreamplotSet:
|
||||
|
||||
def __init__(self, lines, arrows, **kwargs):
|
||||
if kwargs:
|
||||
cbook.warn_deprecated(
|
||||
"3.3",
|
||||
message="Passing arbitrary keyword arguments to StreamplotSet "
|
||||
"is deprecated since %(since) and will become an "
|
||||
"error %(removal)s.")
|
||||
self.lines = lines
|
||||
self.arrows = arrows
|
||||
|
||||
|
||||
# Coordinate definitions
|
||||
# ========================
|
||||
|
||||
class DomainMap:
|
||||
"""
|
||||
Map representing different coordinate systems.
|
||||
|
||||
Coordinate definitions:
|
||||
|
||||
* axes-coordinates goes from 0 to 1 in the domain.
|
||||
* data-coordinates are specified by the input x-y coordinates.
|
||||
* grid-coordinates goes from 0 to N and 0 to M for an N x M grid,
|
||||
where N and M match the shape of the input data.
|
||||
* mask-coordinates goes from 0 to N and 0 to M for an N x M mask,
|
||||
where N and M are user-specified to control the density of streamlines.
|
||||
|
||||
This class also has methods for adding trajectories to the StreamMask.
|
||||
Before adding a trajectory, run `start_trajectory` to keep track of regions
|
||||
crossed by a given trajectory. Later, if you decide the trajectory is bad
|
||||
(e.g., if the trajectory is very short) just call `undo_trajectory`.
|
||||
"""
|
||||
|
||||
def __init__(self, grid, mask):
|
||||
self.grid = grid
|
||||
self.mask = mask
|
||||
# Constants for conversion between grid- and mask-coordinates
|
||||
self.x_grid2mask = (mask.nx - 1) / (grid.nx - 1)
|
||||
self.y_grid2mask = (mask.ny - 1) / (grid.ny - 1)
|
||||
|
||||
self.x_mask2grid = 1. / self.x_grid2mask
|
||||
self.y_mask2grid = 1. / self.y_grid2mask
|
||||
|
||||
self.x_data2grid = 1. / grid.dx
|
||||
self.y_data2grid = 1. / grid.dy
|
||||
|
||||
def grid2mask(self, xi, yi):
|
||||
"""Return nearest space in mask-coords from given grid-coords."""
|
||||
return (int(xi * self.x_grid2mask + 0.5),
|
||||
int(yi * self.y_grid2mask + 0.5))
|
||||
|
||||
def mask2grid(self, xm, ym):
|
||||
return xm * self.x_mask2grid, ym * self.y_mask2grid
|
||||
|
||||
def data2grid(self, xd, yd):
|
||||
return xd * self.x_data2grid, yd * self.y_data2grid
|
||||
|
||||
def grid2data(self, xg, yg):
|
||||
return xg / self.x_data2grid, yg / self.y_data2grid
|
||||
|
||||
def start_trajectory(self, xg, yg):
|
||||
xm, ym = self.grid2mask(xg, yg)
|
||||
self.mask._start_trajectory(xm, ym)
|
||||
|
||||
def reset_start_point(self, xg, yg):
|
||||
xm, ym = self.grid2mask(xg, yg)
|
||||
self.mask._current_xy = (xm, ym)
|
||||
|
||||
def update_trajectory(self, xg, yg):
|
||||
if not self.grid.within_grid(xg, yg):
|
||||
raise InvalidIndexError
|
||||
xm, ym = self.grid2mask(xg, yg)
|
||||
self.mask._update_trajectory(xm, ym)
|
||||
|
||||
def undo_trajectory(self):
|
||||
self.mask._undo_trajectory()
|
||||
|
||||
|
||||
class Grid:
|
||||
"""Grid of data."""
|
||||
def __init__(self, x, y):
|
||||
|
||||
if x.ndim == 1:
|
||||
pass
|
||||
elif x.ndim == 2:
|
||||
x_row = x[0, :]
|
||||
if not np.allclose(x_row, x):
|
||||
raise ValueError("The rows of 'x' must be equal")
|
||||
x = x_row
|
||||
else:
|
||||
raise ValueError("'x' can have at maximum 2 dimensions")
|
||||
|
||||
if y.ndim == 1:
|
||||
pass
|
||||
elif y.ndim == 2:
|
||||
y_col = y[:, 0]
|
||||
if not np.allclose(y_col, y.T):
|
||||
raise ValueError("The columns of 'y' must be equal")
|
||||
y = y_col
|
||||
else:
|
||||
raise ValueError("'y' can have at maximum 2 dimensions")
|
||||
|
||||
self.nx = len(x)
|
||||
self.ny = len(y)
|
||||
|
||||
self.dx = x[1] - x[0]
|
||||
self.dy = y[1] - y[0]
|
||||
|
||||
self.x_origin = x[0]
|
||||
self.y_origin = y[0]
|
||||
|
||||
self.width = x[-1] - x[0]
|
||||
self.height = y[-1] - y[0]
|
||||
|
||||
if not np.allclose(np.diff(x), self.width / (self.nx - 1)):
|
||||
raise ValueError("'x' values must be equally spaced")
|
||||
if not np.allclose(np.diff(y), self.height / (self.ny - 1)):
|
||||
raise ValueError("'y' values must be equally spaced")
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.ny, self.nx
|
||||
|
||||
def within_grid(self, xi, yi):
|
||||
"""Return True if point is a valid index of grid."""
|
||||
# Note that xi/yi can be floats; so, for example, we can't simply check
|
||||
# `xi < self.nx` since *xi* can be `self.nx - 1 < xi < self.nx`
|
||||
return 0 <= xi <= self.nx - 1 and 0 <= yi <= self.ny - 1
|
||||
|
||||
|
||||
class StreamMask:
|
||||
"""
|
||||
Mask to keep track of discrete regions crossed by streamlines.
|
||||
|
||||
The resolution of this grid determines the approximate spacing between
|
||||
trajectories. Streamlines are only allowed to pass through zeroed cells:
|
||||
When a streamline enters a cell, that cell is set to 1, and no new
|
||||
streamlines are allowed to enter.
|
||||
"""
|
||||
|
||||
def __init__(self, density):
|
||||
try:
|
||||
self.nx, self.ny = (30 * np.broadcast_to(density, 2)).astype(int)
|
||||
except ValueError as err:
|
||||
raise ValueError("'density' must be a scalar or be of length "
|
||||
"2") from err
|
||||
if self.nx < 0 or self.ny < 0:
|
||||
raise ValueError("'density' must be positive")
|
||||
self._mask = np.zeros((self.ny, self.nx))
|
||||
self.shape = self._mask.shape
|
||||
|
||||
self._current_xy = None
|
||||
|
||||
def __getitem__(self, args):
|
||||
return self._mask[args]
|
||||
|
||||
def _start_trajectory(self, xm, ym):
|
||||
"""Start recording streamline trajectory"""
|
||||
self._traj = []
|
||||
self._update_trajectory(xm, ym)
|
||||
|
||||
def _undo_trajectory(self):
|
||||
"""Remove current trajectory from mask"""
|
||||
for t in self._traj:
|
||||
self._mask[t] = 0
|
||||
|
||||
def _update_trajectory(self, xm, ym):
|
||||
"""
|
||||
Update current trajectory position in mask.
|
||||
|
||||
If the new position has already been filled, raise `InvalidIndexError`.
|
||||
"""
|
||||
if self._current_xy != (xm, ym):
|
||||
if self[ym, xm] == 0:
|
||||
self._traj.append((ym, xm))
|
||||
self._mask[ym, xm] = 1
|
||||
self._current_xy = (xm, ym)
|
||||
else:
|
||||
raise InvalidIndexError
|
||||
|
||||
|
||||
class InvalidIndexError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class TerminateTrajectory(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# Integrator definitions
|
||||
# =======================
|
||||
|
||||
def get_integrator(u, v, dmap, minlength, maxlength, integration_direction):
|
||||
|
||||
# rescale velocity onto grid-coordinates for integrations.
|
||||
u, v = dmap.data2grid(u, v)
|
||||
|
||||
# speed (path length) will be in axes-coordinates
|
||||
u_ax = u / (dmap.grid.nx - 1)
|
||||
v_ax = v / (dmap.grid.ny - 1)
|
||||
speed = np.ma.sqrt(u_ax ** 2 + v_ax ** 2)
|
||||
|
||||
def forward_time(xi, yi):
|
||||
if not dmap.grid.within_grid(xi, yi):
|
||||
raise OutOfBounds
|
||||
ds_dt = interpgrid(speed, xi, yi)
|
||||
if ds_dt == 0:
|
||||
raise TerminateTrajectory()
|
||||
dt_ds = 1. / ds_dt
|
||||
ui = interpgrid(u, xi, yi)
|
||||
vi = interpgrid(v, xi, yi)
|
||||
return ui * dt_ds, vi * dt_ds
|
||||
|
||||
def backward_time(xi, yi):
|
||||
dxi, dyi = forward_time(xi, yi)
|
||||
return -dxi, -dyi
|
||||
|
||||
def integrate(x0, y0):
|
||||
"""
|
||||
Return x, y grid-coordinates of trajectory based on starting point.
|
||||
|
||||
Integrate both forward and backward in time from starting point in
|
||||
grid coordinates.
|
||||
|
||||
Integration is terminated when a trajectory reaches a domain boundary
|
||||
or when it crosses into an already occupied cell in the StreamMask. The
|
||||
resulting trajectory is None if it is shorter than `minlength`.
|
||||
"""
|
||||
|
||||
stotal, x_traj, y_traj = 0., [], []
|
||||
|
||||
try:
|
||||
dmap.start_trajectory(x0, y0)
|
||||
except InvalidIndexError:
|
||||
return None
|
||||
if integration_direction in ['both', 'backward']:
|
||||
s, xt, yt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength)
|
||||
stotal += s
|
||||
x_traj += xt[::-1]
|
||||
y_traj += yt[::-1]
|
||||
|
||||
if integration_direction in ['both', 'forward']:
|
||||
dmap.reset_start_point(x0, y0)
|
||||
s, xt, yt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength)
|
||||
if len(x_traj) > 0:
|
||||
xt = xt[1:]
|
||||
yt = yt[1:]
|
||||
stotal += s
|
||||
x_traj += xt
|
||||
y_traj += yt
|
||||
|
||||
if stotal > minlength:
|
||||
return x_traj, y_traj
|
||||
else: # reject short trajectories
|
||||
dmap.undo_trajectory()
|
||||
return None
|
||||
|
||||
return integrate
|
||||
|
||||
|
||||
class OutOfBounds(IndexError):
|
||||
pass
|
||||
|
||||
|
||||
def _integrate_rk12(x0, y0, dmap, f, maxlength):
|
||||
"""
|
||||
2nd-order Runge-Kutta algorithm with adaptive step size.
|
||||
|
||||
This method is also referred to as the improved Euler's method, or Heun's
|
||||
method. This method is favored over higher-order methods because:
|
||||
|
||||
1. To get decent looking trajectories and to sample every mask cell
|
||||
on the trajectory we need a small timestep, so a lower order
|
||||
solver doesn't hurt us unless the data is *very* high resolution.
|
||||
In fact, for cases where the user inputs
|
||||
data smaller or of similar grid size to the mask grid, the higher
|
||||
order corrections are negligible because of the very fast linear
|
||||
interpolation used in `interpgrid`.
|
||||
|
||||
2. For high resolution input data (i.e. beyond the mask
|
||||
resolution), we must reduce the timestep. Therefore, an adaptive
|
||||
timestep is more suited to the problem as this would be very hard
|
||||
to judge automatically otherwise.
|
||||
|
||||
This integrator is about 1.5 - 2x as fast as both the RK4 and RK45
|
||||
solvers in most setups on my machine. I would recommend removing the
|
||||
other two to keep things simple.
|
||||
"""
|
||||
# This error is below that needed to match the RK4 integrator. It
|
||||
# is set for visual reasons -- too low and corners start
|
||||
# appearing ugly and jagged. Can be tuned.
|
||||
maxerror = 0.003
|
||||
|
||||
# This limit is important (for all integrators) to avoid the
|
||||
# trajectory skipping some mask cells. We could relax this
|
||||
# condition if we use the code which is commented out below to
|
||||
# increment the location gradually. However, due to the efficient
|
||||
# nature of the interpolation, this doesn't boost speed by much
|
||||
# for quite a bit of complexity.
|
||||
maxds = min(1. / dmap.mask.nx, 1. / dmap.mask.ny, 0.1)
|
||||
|
||||
ds = maxds
|
||||
stotal = 0
|
||||
xi = x0
|
||||
yi = y0
|
||||
xf_traj = []
|
||||
yf_traj = []
|
||||
|
||||
while True:
|
||||
try:
|
||||
if dmap.grid.within_grid(xi, yi):
|
||||
xf_traj.append(xi)
|
||||
yf_traj.append(yi)
|
||||
else:
|
||||
raise OutOfBounds
|
||||
|
||||
# Compute the two intermediate gradients.
|
||||
# f should raise OutOfBounds if the locations given are
|
||||
# outside the grid.
|
||||
k1x, k1y = f(xi, yi)
|
||||
k2x, k2y = f(xi + ds * k1x, yi + ds * k1y)
|
||||
|
||||
except OutOfBounds:
|
||||
# Out of the domain during this step.
|
||||
# Take an Euler step to the boundary to improve neatness
|
||||
# unless the trajectory is currently empty.
|
||||
if xf_traj:
|
||||
ds, xf_traj, yf_traj = _euler_step(xf_traj, yf_traj,
|
||||
dmap, f)
|
||||
stotal += ds
|
||||
break
|
||||
except TerminateTrajectory:
|
||||
break
|
||||
|
||||
dx1 = ds * k1x
|
||||
dy1 = ds * k1y
|
||||
dx2 = ds * 0.5 * (k1x + k2x)
|
||||
dy2 = ds * 0.5 * (k1y + k2y)
|
||||
|
||||
nx, ny = dmap.grid.shape
|
||||
# Error is normalized to the axes coordinates
|
||||
error = np.hypot((dx2 - dx1) / (nx - 1), (dy2 - dy1) / (ny - 1))
|
||||
|
||||
# Only save step if within error tolerance
|
||||
if error < maxerror:
|
||||
xi += dx2
|
||||
yi += dy2
|
||||
try:
|
||||
dmap.update_trajectory(xi, yi)
|
||||
except InvalidIndexError:
|
||||
break
|
||||
if stotal + ds > maxlength:
|
||||
break
|
||||
stotal += ds
|
||||
|
||||
# recalculate stepsize based on step error
|
||||
if error == 0:
|
||||
ds = maxds
|
||||
else:
|
||||
ds = min(maxds, 0.85 * ds * (maxerror / error) ** 0.5)
|
||||
|
||||
return stotal, xf_traj, yf_traj
|
||||
|
||||
|
||||
def _euler_step(xf_traj, yf_traj, dmap, f):
|
||||
"""Simple Euler integration step that extends streamline to boundary."""
|
||||
ny, nx = dmap.grid.shape
|
||||
xi = xf_traj[-1]
|
||||
yi = yf_traj[-1]
|
||||
cx, cy = f(xi, yi)
|
||||
if cx == 0:
|
||||
dsx = np.inf
|
||||
elif cx < 0:
|
||||
dsx = xi / -cx
|
||||
else:
|
||||
dsx = (nx - 1 - xi) / cx
|
||||
if cy == 0:
|
||||
dsy = np.inf
|
||||
elif cy < 0:
|
||||
dsy = yi / -cy
|
||||
else:
|
||||
dsy = (ny - 1 - yi) / cy
|
||||
ds = min(dsx, dsy)
|
||||
xf_traj.append(xi + cx * ds)
|
||||
yf_traj.append(yi + cy * ds)
|
||||
return ds, xf_traj, yf_traj
|
||||
|
||||
|
||||
# Utility functions
|
||||
# ========================
|
||||
|
||||
def interpgrid(a, xi, yi):
|
||||
"""Fast 2D, linear interpolation on an integer grid"""
|
||||
|
||||
Ny, Nx = np.shape(a)
|
||||
if isinstance(xi, np.ndarray):
|
||||
x = xi.astype(int)
|
||||
y = yi.astype(int)
|
||||
# Check that xn, yn don't exceed max index
|
||||
xn = np.clip(x + 1, 0, Nx - 1)
|
||||
yn = np.clip(y + 1, 0, Ny - 1)
|
||||
else:
|
||||
x = int(xi)
|
||||
y = int(yi)
|
||||
# conditional is faster than clipping for integers
|
||||
if x == (Nx - 1):
|
||||
xn = x
|
||||
else:
|
||||
xn = x + 1
|
||||
if y == (Ny - 1):
|
||||
yn = y
|
||||
else:
|
||||
yn = y + 1
|
||||
|
||||
a00 = a[y, x]
|
||||
a01 = a[y, xn]
|
||||
a10 = a[yn, x]
|
||||
a11 = a[yn, xn]
|
||||
xt = xi - x
|
||||
yt = yi - y
|
||||
a0 = a00 * (1 - xt) + a01 * xt
|
||||
a1 = a10 * (1 - xt) + a11 * xt
|
||||
ai = a0 * (1 - yt) + a1 * yt
|
||||
|
||||
if not isinstance(xi, np.ndarray):
|
||||
if np.ma.is_masked(ai):
|
||||
raise TerminateTrajectory
|
||||
|
||||
return ai
|
||||
|
||||
|
||||
def _gen_starting_points(shape):
|
||||
"""
|
||||
Yield starting points for streamlines.
|
||||
|
||||
Trying points on the boundary first gives higher quality streamlines.
|
||||
This algorithm starts with a point on the mask corner and spirals inward.
|
||||
This algorithm is inefficient, but fast compared to rest of streamplot.
|
||||
"""
|
||||
ny, nx = shape
|
||||
xfirst = 0
|
||||
yfirst = 1
|
||||
xlast = nx - 1
|
||||
ylast = ny - 1
|
||||
x, y = 0, 0
|
||||
direction = 'right'
|
||||
for i in range(nx * ny):
|
||||
yield x, y
|
||||
|
||||
if direction == 'right':
|
||||
x += 1
|
||||
if x >= xlast:
|
||||
xlast -= 1
|
||||
direction = 'up'
|
||||
elif direction == 'up':
|
||||
y += 1
|
||||
if y >= ylast:
|
||||
ylast -= 1
|
||||
direction = 'left'
|
||||
elif direction == 'left':
|
||||
x -= 1
|
||||
if x <= xfirst:
|
||||
xfirst += 1
|
||||
direction = 'down'
|
||||
elif direction == 'down':
|
||||
y -= 1
|
||||
if y <= yfirst:
|
||||
yfirst += 1
|
||||
direction = 'right'
|
Loading…
Add table
Add a link
Reference in a new issue