cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',

minlength=0.1, transform=None, zorder=2, start_points=None):

"""Draws streamlines of a vector flow.

grid = Grid(x, y)

mask = StreamMask(density)

dmap = DomainMap(grid, mask)

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

use_multicolor_lines = isinstance(color, np.ndarray)

if use_multicolor_lines:

if color.shape != grid.shape:

msg = "If 'color' is given, must have the shape of 'Grid(x,y)'"

raise ValueError(msg)

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:

msg = "If 'linewidth' is given, must have the shape of 'Grid(x,y)'"

raise ValueError(msg)

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

msg = "'u' and 'v' must be of shape 'Grid(x,y)'"

raise ValueError(msg)

u = np.ma.masked_invalid(u)

v = np.ma.masked_invalid(v)

integrate = get_integrator(u, v, dmap, minlength)

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=np.float).copy()

# Check if start_points are outside the data boundaries

for xs, ys in sp2:

if (xs < grid.x_origin or xs > grid.x_origin + grid.width

or ys < grid.y_origin or 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.sqrt(np.diff(tx) ** 2 + np.diff(ty) ** 2))

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,

margins=False,

**arrow_kw)

axes.add_patch(p)

arrows.append(p)

lc = mcollections.LineCollection(streamlines,

transform=transform,

margins=False,

**line_kw)

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, margins=False)

stream_container = StreamplotSet(lc, ac)

def __init__(self, lines, arrows, **kwargs):

self.lines = lines

"""Map representing different coordinate systems.

def __init__(self, grid, mask):

self.grid = grid

self.mask = mask

# Constants for conversion between grid- and mask-coordinates

self.x_grid2mask = float(mask.nx - 1) / grid.nx

self.y_grid2mask = float(mask.ny - 1) / grid.ny

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

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

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

"""Mask to keep track of discrete regions crossed by streamlines.

def __init__(self, density):

if np.isscalar(density):

if density <= 0:

raise ValueError("If a scalar, 'density' must be positive")

self.nx = self.ny = int(30 * density)

else:

if len(density) != 2:

raise ValueError("'density' can have at maximum 2 dimensions")

self.nx = int(30 * density[0])

self.ny = int(30 * density[1])

self._mask = np.zeros((self.ny, self.nx))

self.shape = self._mask.shape

self._current_xy = None

def __getitem__(self, *args):

return self._mask.__getitem__(*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.__setitem__(t, 0)

def _update_trajectory(self, xm, ym):

"""Update current trajectory position in mask.

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:

# 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

v_ax = v / dmap.grid.ny

speed = np.ma.sqrt(u_ax ** 2 + v_ax ** 2)

def forward_time(xi, yi):

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.

try:

dmap.start_trajectory(x0, y0)

except InvalidIndexError:

return None

sf, xf_traj, yf_traj = _integrate_rk12(x0, y0, dmap, forward_time)

dmap.reset_start_point(x0, y0)

sb, xb_traj, yb_traj = _integrate_rk12(x0, y0, dmap, backward_time)

# combine forward and backward trajectories

stotal = sf + sb

x_traj = xb_traj[::-1] + xf_traj[1:]

y_traj = yb_traj[::-1] + yf_traj[1:]

if stotal > minlength:

return x_traj, y_traj

else: # reject short trajectories

dmap.undo_trajectory()

return None

"""2nd-order Runge-Kutta algorithm with adaptive step size.

# 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 dmap.grid.within_grid(xi, yi):

xf_traj.append(xi)

yf_traj.append(yi)

try:

k1x, k1y = f(xi, yi)

k2x, k2y = f(xi + ds * k1x,

yi + ds * k1y)

except IndexError:

# Out of the domain on one of the intermediate integration steps.

# Take an Euler step to the boundary to improve neatness.

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.sqrt(((dx2 - dx1) / nx) ** 2 + ((dy2 - dy1) / ny) ** 2)

# 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) > 2:

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)

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

"""Fast 2D, linear interpolation on an integer grid"""

Ny, Nx = np.shape(a)

if isinstance(xi, np.ndarray):

x = xi.astype(np.int)

y = yi.astype(np.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 = np.int(xi)

y = np.int(yi)

# conditional is faster than clipping for integers

if x == (Nx - 2):

xn = x

else:

xn = x + 1

if y == (Ny - 2):

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

"""Yield starting points for streamlines.

ny, nx = shape

xfirst = 0

yfirst = 1

xlast = nx - 1

ylast = ny - 1

x, y = 0, 0

i = 0

direction = 'right'

for i in xrange(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