Coverage for /usr/lib/python3/dist-packages/matplotlib/streamplot.py: 12%

1""" 

2Streamline plotting for 2D vector fields. 

3 

4""" 

5 

6import numpy as np 

7 

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 

13 

14 

15__all__ = ['streamplot'] 

16 

17 

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. 

25 

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 : 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 : Nx2 array 

61 Coordinates of starting points for the streamlines in data coordinates 

62 (the same coordinates as the *x* and *y* arrays). 

63 zorder : int 

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. 

76 

77 Returns 

78 ------- 

79 StreamplotSet 

80 Container object with attributes 

81 

82 - ``lines``: `.LineCollection` of streamlines 

83 

84 - ``arrows``: `.PatchCollection` containing `.FancyArrowPatch` 

85 objects representing the arrows half-way along streamlines. 

86 

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) 

94 

95 if zorder is None: 

96 zorder = mlines.Line2D.zorder 

97 

98 # default to data coordinates 

99 if transform is None: 

100 transform = axes.transData 

101 

102 if color is None: 

103 color = axes._get_lines.get_next_color() 

104 

105 if linewidth is None: 

106 linewidth = mpl.rcParams['lines.linewidth'] 

107 

108 line_kw = {} 

109 arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10 * arrowsize) 

110 

111 _api.check_in_list(['both', 'forward', 'backward'], 

112 integration_direction=integration_direction) 

113 

114 if integration_direction == 'both': 

115 maxlength /= 2. 

116 

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 

127 

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 

136 

137 line_kw['zorder'] = zorder 

138 arrow_kw['zorder'] = zorder 

139 

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

143 

144 u = np.ma.masked_invalid(u) 

145 v = np.ma.masked_invalid(v) 

146 

147 integrate = _get_integrator(u, v, dmap, minlength, maxlength, 

148 integration_direction) 

149 

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

160 

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("Starting point ({}, {}) outside of data " 

166 "boundaries".format(xs, ys)) 

167 

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 

173 

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) 

182 

183 t = integrate(xg, yg, broken_streamlines) 

184 if t is not None: 

185 trajectories.append(t) 

186 

187 if use_multicolor_lines: 

188 if norm is None: 

189 norm = mcolors.Normalize(color.min(), color.max()) 

190 cmap = cm._ensure_cmap(cmap) 

191 

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 

200 

201 points = np.transpose([tx, ty]).reshape(-1, 1, 2) 

202 streamlines.extend(np.hstack([points[:-1], points[1:]])) 

203 

204 # Add arrows halfway along each trajectory. 

205 s = np.cumsum(np.hypot(np.diff(tx), np.diff(ty))) 

206 n = np.searchsorted(s, s[-1] / 2.) 

207 arrow_tail = (tx[n], ty[n]) 

208 arrow_head = (np.mean(tx[n:n + 2]), np.mean(ty[n:n + 2])) 

209 

210 if isinstance(linewidth, np.ndarray): 

211 line_widths = interpgrid(linewidth, tgx, tgy)[:-1] 

212 line_kw['linewidth'].extend(line_widths) 

213 arrow_kw['linewidth'] = line_widths[n] 

214 

215 if use_multicolor_lines: 

216 color_values = interpgrid(color, tgx, tgy)[:-1] 

217 line_colors.append(color_values) 

218 arrow_kw['color'] = cmap(norm(color_values[n])) 

219 

220 p = patches.FancyArrowPatch( 

221 arrow_tail, arrow_head, transform=transform, **arrow_kw) 

222 arrows.append(p) 

223 

224 lc = mcollections.LineCollection( 

225 streamlines, transform=transform, **line_kw) 

226 lc.sticky_edges.x[:] = [grid.x_origin, grid.x_origin + grid.width] 

227 lc.sticky_edges.y[:] = [grid.y_origin, grid.y_origin + grid.height] 

228 if use_multicolor_lines: 

229 lc.set_array(np.ma.hstack(line_colors)) 

230 lc.set_cmap(cmap) 

231 lc.set_norm(norm) 

232 axes.add_collection(lc) 

233 

234 ac = mcollections.PatchCollection(arrows) 

235 # Adding the collection itself is broken; see #2341. 

236 for p in arrows: 

237 axes.add_patch(p) 

238 

239 axes.autoscale_view() 

240 stream_container = StreamplotSet(lc, ac) 

241 return stream_container 

242 

243 

244class StreamplotSet: 

245 

246 def __init__(self, lines, arrows): 

247 self.lines = lines 

248 self.arrows = arrows 

249 

250 

251# Coordinate definitions 

252# ======================== 

253 

254class DomainMap: 

255 """ 

256 Map representing different coordinate systems. 

257 

258 Coordinate definitions: 

259 

260 * axes-coordinates goes from 0 to 1 in the domain. 

261 * data-coordinates are specified by the input x-y coordinates. 

262 * grid-coordinates goes from 0 to N and 0 to M for an N x M grid, 

263 where N and M match the shape of the input data. 

264 * mask-coordinates goes from 0 to N and 0 to M for an N x M mask, 

265 where N and M are user-specified to control the density of streamlines. 

266 

267 This class also has methods for adding trajectories to the StreamMask. 

268 Before adding a trajectory, run `start_trajectory` to keep track of regions 

269 crossed by a given trajectory. Later, if you decide the trajectory is bad 

270 (e.g., if the trajectory is very short) just call `undo_trajectory`. 

271 """ 

272 

273 def __init__(self, grid, mask): 

274 self.grid = grid 

275 self.mask = mask 

276 # Constants for conversion between grid- and mask-coordinates 

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

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

279 

280 self.x_mask2grid = 1. / self.x_grid2mask 

281 self.y_mask2grid = 1. / self.y_grid2mask 

282 

283 self.x_data2grid = 1. / grid.dx 

284 self.y_data2grid = 1. / grid.dy 

285 

286 def grid2mask(self, xi, yi): 

287 """Return nearest space in mask-coords from given grid-coords.""" 

288 return (int(xi * self.x_grid2mask + 0.5), 

289 int(yi * self.y_grid2mask + 0.5)) 

290 

291 def mask2grid(self, xm, ym): 

292 return xm * self.x_mask2grid, ym * self.y_mask2grid 

293 

294 def data2grid(self, xd, yd): 

295 return xd * self.x_data2grid, yd * self.y_data2grid 

296 

297 def grid2data(self, xg, yg): 

298 return xg / self.x_data2grid, yg / self.y_data2grid 

299 

300 def start_trajectory(self, xg, yg, broken_streamlines=True): 

301 xm, ym = self.grid2mask(xg, yg) 

302 self.mask._start_trajectory(xm, ym, broken_streamlines) 

303 

304 def reset_start_point(self, xg, yg): 

305 xm, ym = self.grid2mask(xg, yg) 

306 self.mask._current_xy = (xm, ym) 

307 

308 def update_trajectory(self, xg, yg, broken_streamlines=True): 

309 if not self.grid.within_grid(xg, yg): 

310 raise InvalidIndexError 

311 xm, ym = self.grid2mask(xg, yg) 

312 self.mask._update_trajectory(xm, ym, broken_streamlines) 

313 

314 def undo_trajectory(self): 

315 self.mask._undo_trajectory() 

316 

317 

318class Grid: 

319 """Grid of data.""" 

320 def __init__(self, x, y): 

321 

322 if np.ndim(x) == 1: 

323 pass 

324 elif np.ndim(x) == 2: 

325 x_row = x[0] 

326 if not np.allclose(x_row, x): 

327 raise ValueError("The rows of 'x' must be equal") 

328 x = x_row 

329 else: 

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

331 

332 if np.ndim(y) == 1: 

333 pass 

334 elif np.ndim(y) == 2: 

335 yt = np.transpose(y) # Also works for nested lists. 

336 y_col = yt[0] 

337 if not np.allclose(y_col, yt): 

338 raise ValueError("The columns of 'y' must be equal") 

339 y = y_col 

340 else: 

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

342 

343 if not (np.diff(x) > 0).all(): 

344 raise ValueError("'x' must be strictly increasing") 

345 if not (np.diff(y) > 0).all(): 

346 raise ValueError("'y' must be strictly increasing") 

347 

348 self.nx = len(x) 

349 self.ny = len(y) 

350 

351 self.dx = x[1] - x[0] 

352 self.dy = y[1] - y[0] 

353 

354 self.x_origin = x[0] 

355 self.y_origin = y[0] 

356 

357 self.width = x[-1] - x[0] 

358 self.height = y[-1] - y[0] 

359 

360 if not np.allclose(np.diff(x), self.width / (self.nx - 1)): 

361 raise ValueError("'x' values must be equally spaced") 

362 if not np.allclose(np.diff(y), self.height / (self.ny - 1)): 

363 raise ValueError("'y' values must be equally spaced") 

364 

365 @property 

366 def shape(self): 

367 return self.ny, self.nx 

368 

369 def within_grid(self, xi, yi): 

370 """Return whether (*xi*, *yi*) is a valid index of the grid.""" 

371 # Note that xi/yi can be floats; so, for example, we can't simply check 

372 # `xi < self.nx` since *xi* can be `self.nx - 1 < xi < self.nx` 

373 return 0 <= xi <= self.nx - 1 and 0 <= yi <= self.ny - 1 

374 

375 

376class StreamMask: 

377 """ 

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

379 

380 The resolution of this grid determines the approximate spacing between 

381 trajectories. Streamlines are only allowed to pass through zeroed cells: 

382 When a streamline enters a cell, that cell is set to 1, and no new 

383 streamlines are allowed to enter. 

384 """ 

385 

386 def __init__(self, density): 

387 try: 

388 self.nx, self.ny = (30 * np.broadcast_to(density, 2)).astype(int) 

389 except ValueError as err: 

390 raise ValueError("'density' must be a scalar or be of length " 

391 "2") from err 

392 if self.nx < 0 or self.ny < 0: 

393 raise ValueError("'density' must be positive") 

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

395 self.shape = self._mask.shape 

396 

397 self._current_xy = None 

398 

399 def __getitem__(self, args): 

400 return self._mask[args] 

401 

402 def _start_trajectory(self, xm, ym, broken_streamlines=True): 

403 """Start recording streamline trajectory""" 

404 self._traj = [] 

405 self._update_trajectory(xm, ym, broken_streamlines) 

406 

407 def _undo_trajectory(self): 

408 """Remove current trajectory from mask""" 

409 for t in self._traj: 

410 self._mask[t] = 0 

411 

412 def _update_trajectory(self, xm, ym, broken_streamlines=True): 

413 """ 

414 Update current trajectory position in mask. 

415 

416 If the new position has already been filled, raise `InvalidIndexError`. 

417 """ 

418 if self._current_xy != (xm, ym): 

419 if self[ym, xm] == 0: 

420 self._traj.append((ym, xm)) 

421 self._mask[ym, xm] = 1 

422 self._current_xy = (xm, ym) 

423 else: 

424 if broken_streamlines: 

425 raise InvalidIndexError 

426 else: 

427 pass 

428 

429 

430class InvalidIndexError(Exception): 

431 pass 

432 

433 

434class TerminateTrajectory(Exception): 

435 pass 

436 

437 

438# Integrator definitions 

439# ======================= 

440 

441def _get_integrator(u, v, dmap, minlength, maxlength, integration_direction): 

442 

443 # rescale velocity onto grid-coordinates for integrations. 

444 u, v = dmap.data2grid(u, v) 

445 

446 # speed (path length) will be in axes-coordinates 

447 u_ax = u / (dmap.grid.nx - 1) 

448 v_ax = v / (dmap.grid.ny - 1) 

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

450 

451 def forward_time(xi, yi): 

452 if not dmap.grid.within_grid(xi, yi): 

453 raise OutOfBounds 

454 ds_dt = interpgrid(speed, xi, yi) 

455 if ds_dt == 0: 

456 raise TerminateTrajectory() 

457 dt_ds = 1. / ds_dt 

458 ui = interpgrid(u, xi, yi) 

459 vi = interpgrid(v, xi, yi) 

460 return ui * dt_ds, vi * dt_ds 

461 

462 def backward_time(xi, yi): 

463 dxi, dyi = forward_time(xi, yi) 

464 return -dxi, -dyi 

465 

466 def integrate(x0, y0, broken_streamlines=True): 

467 """ 

468 Return x, y grid-coordinates of trajectory based on starting point. 

469 

470 Integrate both forward and backward in time from starting point in 

471 grid coordinates. 

472 

473 Integration is terminated when a trajectory reaches a domain boundary 

474 or when it crosses into an already occupied cell in the StreamMask. The 

475 resulting trajectory is None if it is shorter than `minlength`. 

476 """ 

477 

478 stotal, xy_traj = 0., [] 

479 

480 try: 

481 dmap.start_trajectory(x0, y0, broken_streamlines) 

482 except InvalidIndexError: 

483 return None 

484 if integration_direction in ['both', 'backward']: 

485 s, xyt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength, 

486 broken_streamlines) 

487 stotal += s 

488 xy_traj += xyt[::-1] 

489 

490 if integration_direction in ['both', 'forward']: 

491 dmap.reset_start_point(x0, y0) 

492 s, xyt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength, 

493 broken_streamlines) 

494 stotal += s 

495 xy_traj += xyt[1:] 

496 

497 if stotal > minlength: 

498 return np.broadcast_arrays(xy_traj, np.empty((1, 2)))[0] 

499 else: # reject short trajectories 

500 dmap.undo_trajectory() 

501 return None 

502 

503 return integrate 

504 

505 

506@_api.deprecated("3.5") 

507def get_integrator(u, v, dmap, minlength, maxlength, integration_direction): 

508 xy_traj = _get_integrator( 

509 u, v, dmap, minlength, maxlength, integration_direction) 

510 return (None if xy_traj is None 

511 else ([], []) if not len(xy_traj) 

512 else [*zip(*xy_traj)]) 

513 

514 

515class OutOfBounds(IndexError): 

516 pass 

517 

518 

519def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True): 

520 """ 

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

522 

523 This method is also referred to as the improved Euler's method, or Heun's 

524 method. This method is favored over higher-order methods because: 

525 

526 1. To get decent looking trajectories and to sample every mask cell 

527 on the trajectory we need a small timestep, so a lower order 

528 solver doesn't hurt us unless the data is *very* high resolution. 

529 In fact, for cases where the user inputs 

530 data smaller or of similar grid size to the mask grid, the higher 

531 order corrections are negligible because of the very fast linear 

532 interpolation used in `interpgrid`. 

533 

534 2. For high resolution input data (i.e. beyond the mask 

535 resolution), we must reduce the timestep. Therefore, an adaptive 

536 timestep is more suited to the problem as this would be very hard 

537 to judge automatically otherwise. 

538 

539 This integrator is about 1.5 - 2x as fast as RK4 and RK45 solvers (using 

540 similar Python implementations) in most setups. 

541 """ 

542 # This error is below that needed to match the RK4 integrator. It 

543 # is set for visual reasons -- too low and corners start 

544 # appearing ugly and jagged. Can be tuned. 

545 maxerror = 0.003 

546 

547 # This limit is important (for all integrators) to avoid the 

548 # trajectory skipping some mask cells. We could relax this 

549 # condition if we use the code which is commented out below to 

550 # increment the location gradually. However, due to the efficient 

551 # nature of the interpolation, this doesn't boost speed by much 

552 # for quite a bit of complexity. 

553 maxds = min(1. / dmap.mask.nx, 1. / dmap.mask.ny, 0.1) 

554 

555 ds = maxds 

556 stotal = 0 

557 xi = x0 

558 yi = y0 

559 xyf_traj = [] 

560 

561 while True: 

562 try: 

563 if dmap.grid.within_grid(xi, yi): 

564 xyf_traj.append((xi, yi)) 

565 else: 

566 raise OutOfBounds 

567 

568 # Compute the two intermediate gradients. 

569 # f should raise OutOfBounds if the locations given are 

570 # outside the grid. 

571 k1x, k1y = f(xi, yi) 

572 k2x, k2y = f(xi + ds * k1x, yi + ds * k1y) 

573 

574 except OutOfBounds: 

575 # Out of the domain during this step. 

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

577 # unless the trajectory is currently empty. 

578 if xyf_traj: 

579 ds, xyf_traj = _euler_step(xyf_traj, dmap, f) 

580 stotal += ds 

581 break 

582 except TerminateTrajectory: 

583 break 

584 

585 dx1 = ds * k1x 

586 dy1 = ds * k1y 

587 dx2 = ds * 0.5 * (k1x + k2x) 

588 dy2 = ds * 0.5 * (k1y + k2y) 

589 

590 ny, nx = dmap.grid.shape 

591 # Error is normalized to the axes coordinates 

592 error = np.hypot((dx2 - dx1) / (nx - 1), (dy2 - dy1) / (ny - 1)) 

593 

594 # Only save step if within error tolerance 

595 if error < maxerror: 

596 xi += dx2 

597 yi += dy2 

598 try: 

599 dmap.update_trajectory(xi, yi, broken_streamlines) 

600 except InvalidIndexError: 

601 break 

602 if stotal + ds > maxlength: 

603 break 

604 stotal += ds 

605 

606 # recalculate stepsize based on step error 

607 if error == 0: 

608 ds = maxds 

609 else: 

610 ds = min(maxds, 0.85 * ds * (maxerror / error) ** 0.5) 

611 

612 return stotal, xyf_traj 

613 

614 

615def _euler_step(xyf_traj, dmap, f): 

616 """Simple Euler integration step that extends streamline to boundary.""" 

617 ny, nx = dmap.grid.shape 

618 xi, yi = xyf_traj[-1] 

619 cx, cy = f(xi, yi) 

620 if cx == 0: 

621 dsx = np.inf 

622 elif cx < 0: 

623 dsx = xi / -cx 

624 else: 

625 dsx = (nx - 1 - xi) / cx 

626 if cy == 0: 

627 dsy = np.inf 

628 elif cy < 0: 

629 dsy = yi / -cy 

630 else: 

631 dsy = (ny - 1 - yi) / cy 

632 ds = min(dsx, dsy) 

633 xyf_traj.append((xi + cx * ds, yi + cy * ds)) 

634 return ds, xyf_traj 

635 

636 

637# Utility functions 

638# ======================== 

639 

640def interpgrid(a, xi, yi): 

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

642 

643 Ny, Nx = np.shape(a) 

644 if isinstance(xi, np.ndarray): 

645 x = xi.astype(int) 

646 y = yi.astype(int) 

647 # Check that xn, yn don't exceed max index 

648 xn = np.clip(x + 1, 0, Nx - 1) 

649 yn = np.clip(y + 1, 0, Ny - 1) 

650 else: 

651 x = int(xi) 

652 y = int(yi) 

653 # conditional is faster than clipping for integers 

654 if x == (Nx - 1): 

655 xn = x 

656 else: 

657 xn = x + 1 

658 if y == (Ny - 1): 

659 yn = y 

660 else: 

661 yn = y + 1 

662 

663 a00 = a[y, x] 

664 a01 = a[y, xn] 

665 a10 = a[yn, x] 

666 a11 = a[yn, xn] 

667 xt = xi - x 

668 yt = yi - y 

669 a0 = a00 * (1 - xt) + a01 * xt 

670 a1 = a10 * (1 - xt) + a11 * xt 

671 ai = a0 * (1 - yt) + a1 * yt 

672 

673 if not isinstance(xi, np.ndarray): 

674 if np.ma.is_masked(ai): 

675 raise TerminateTrajectory 

676 

677 return ai 

678 

679 

680def _gen_starting_points(shape): 

681 """ 

682 Yield starting points for streamlines. 

683 

684 Trying points on the boundary first gives higher quality streamlines. 

685 This algorithm starts with a point on the mask corner and spirals inward. 

686 This algorithm is inefficient, but fast compared to rest of streamplot. 

687 """ 

688 ny, nx = shape 

689 xfirst = 0 

690 yfirst = 1 

691 xlast = nx - 1 

692 ylast = ny - 1 

693 x, y = 0, 0 

694 direction = 'right' 

695 for i in range(nx * ny): 

696 yield x, y 

697 

698 if direction == 'right': 

699 x += 1 

700 if x >= xlast: 

701 xlast -= 1 

702 direction = 'up' 

703 elif direction == 'up': 

704 y += 1 

705 if y >= ylast: 

706 ylast -= 1 

707 direction = 'left' 

708 elif direction == 'left': 

709 x -= 1 

710 if x <= xfirst: 

711 xfirst += 1 

712 direction = 'down' 

713 elif direction == 'down': 

714 y -= 1 

715 if y <= yfirst: 

716 yfirst += 1 

717 direction = 'right'