(u, v, dmap, minlength, maxlength, integration_direction)
| 409 | #======================== |
| 410 | |
| 411 | def get_integrator(u, v, dmap, minlength, maxlength, integration_direction): |
| 412 | |
| 413 | # rescale velocity onto grid-coordinates for integrations. |
| 414 | u, v = dmap.data2grid(u, v) |
| 415 | |
| 416 | # speed (path length) will be in axes-coordinates |
| 417 | u_ax = u / dmap.grid.nx |
| 418 | v_ax = v / dmap.grid.ny |
| 419 | speed = np.ma.sqrt(u_ax ** 2 + v_ax ** 2) |
| 420 | |
| 421 | def forward_time(xi, yi): |
| 422 | ds_dt = interpgrid(speed, xi, yi) |
| 423 | if ds_dt == 0: |
| 424 | raise TerminateTrajectory() |
| 425 | dt_ds = 1. / ds_dt |
| 426 | ui = interpgrid(u, xi, yi) |
| 427 | vi = interpgrid(v, xi, yi) |
| 428 | return ui * dt_ds, vi * dt_ds |
| 429 | |
| 430 | def backward_time(xi, yi): |
| 431 | dxi, dyi = forward_time(xi, yi) |
| 432 | return -dxi, -dyi |
| 433 | |
| 434 | def integrate(x0, y0): |
| 435 | """Return x, y grid-coordinates of trajectory based on starting point. |
| 436 | |
| 437 | Integrate both forward and backward in time from starting point in |
| 438 | grid coordinates. |
| 439 | |
| 440 | Integration is terminated when a trajectory reaches a domain boundary |
| 441 | or when it crosses into an already occupied cell in the StreamMask. The |
| 442 | resulting trajectory is None if it is shorter than `minlength`. |
| 443 | """ |
| 444 | |
| 445 | stotal, x_traj, y_traj = 0., [], [] |
| 446 | |
| 447 | try: |
| 448 | dmap.start_trajectory(x0, y0) |
| 449 | except InvalidIndexError: |
| 450 | return None |
| 451 | if integration_direction in ['both', 'backward']: |
| 452 | s, xt, yt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength) |
| 453 | stotal += s |
| 454 | x_traj += xt[::-1] |
| 455 | y_traj += yt[::-1] |
| 456 | |
| 457 | if integration_direction in ['both', 'forward']: |
| 458 | dmap.reset_start_point(x0, y0) |
| 459 | s, xt, yt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength) |
| 460 | if len(x_traj) > 0: |
| 461 | xt = xt[1:] |
| 462 | yt = yt[1:] |
| 463 | stotal += s |
| 464 | x_traj += xt |
| 465 | y_traj += yt |
| 466 | |
| 467 | if stotal > minlength: |
| 468 | return x_traj, y_traj |
no test coverage detected