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 stre
| 425 | |
| 426 | |
| 427 | class StreamMask: |
| 428 | """ |
| 429 | Mask to keep track of discrete regions crossed by streamlines. |
| 430 | |
| 431 | The resolution of this grid determines the approximate spacing between |
| 432 | trajectories. Streamlines are only allowed to pass through zeroed cells: |
| 433 | When a streamline enters a cell, that cell is set to 1, and no new |
| 434 | streamlines are allowed to enter. |
| 435 | """ |
| 436 | |
| 437 | def __init__(self, density): |
| 438 | try: |
| 439 | self.nx, self.ny = (30 * np.broadcast_to(density, 2)).astype(int) |
| 440 | except ValueError as err: |
| 441 | raise ValueError("'density' must be a scalar or be of length " |
| 442 | "2") from err |
| 443 | if self.nx < 0 or self.ny < 0: |
| 444 | raise ValueError("'density' must be positive") |
| 445 | self._mask = np.zeros((self.ny, self.nx)) |
| 446 | self.shape = self._mask.shape |
| 447 | |
| 448 | self._current_xy = None |
| 449 | |
| 450 | def __getitem__(self, args): |
| 451 | return self._mask[args] |
| 452 | |
| 453 | def _start_trajectory(self, xm, ym, broken_streamlines=True): |
| 454 | """Start recording streamline trajectory""" |
| 455 | self._traj = [] |
| 456 | self._update_trajectory(xm, ym, broken_streamlines) |
| 457 | |
| 458 | def _undo_trajectory(self): |
| 459 | """Remove current trajectory from mask""" |
| 460 | for t in self._traj: |
| 461 | self._mask[t] = 0 |
| 462 | |
| 463 | def _update_trajectory(self, xm, ym, broken_streamlines=True): |
| 464 | """ |
| 465 | Update current trajectory position in mask. |
| 466 | |
| 467 | If the new position has already been filled, raise `InvalidIndexError`. |
| 468 | """ |
| 469 | if self._current_xy != (xm, ym): |
| 470 | if self[ym, xm] == 0: |
| 471 | self._traj.append((ym, xm)) |
| 472 | self._mask[ym, xm] = 1 |
| 473 | self._current_xy = (xm, ym) |
| 474 | else: |
| 475 | if broken_streamlines: |
| 476 | raise InvalidIndexError |
| 477 | else: |
| 478 | pass |
| 479 | |
| 480 | |
| 481 | class InvalidIndexError(Exception): |
no outgoing calls
no test coverage detected
searching dependent graphs…