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 streamlin
| 347 | |
| 348 | |
| 349 | class StreamMask(object): |
| 350 | """Mask to keep track of discrete regions crossed by streamlines. |
| 351 | |
| 352 | The resolution of this grid determines the approximate spacing between |
| 353 | trajectories. Streamlines are only allowed to pass through zeroed cells: |
| 354 | When a streamline enters a cell, that cell is set to 1, and no new |
| 355 | streamlines are allowed to enter. |
| 356 | """ |
| 357 | |
| 358 | def __init__(self, density): |
| 359 | if np.isscalar(density): |
| 360 | if density <= 0: |
| 361 | raise ValueError("If a scalar, 'density' must be positive") |
| 362 | self.nx = self.ny = int(30 * density) |
| 363 | else: |
| 364 | if len(density) != 2: |
| 365 | raise ValueError("'density' can have at maximum 2 dimensions") |
| 366 | self.nx = int(30 * density[0]) |
| 367 | self.ny = int(30 * density[1]) |
| 368 | self._mask = np.zeros((self.ny, self.nx)) |
| 369 | self.shape = self._mask.shape |
| 370 | |
| 371 | self._current_xy = None |
| 372 | |
| 373 | def __getitem__(self, *args): |
| 374 | return self._mask.__getitem__(*args) |
| 375 | |
| 376 | def _start_trajectory(self, xm, ym): |
| 377 | """Start recording streamline trajectory""" |
| 378 | self._traj = [] |
| 379 | self._update_trajectory(xm, ym) |
| 380 | |
| 381 | def _undo_trajectory(self): |
| 382 | """Remove current trajectory from mask""" |
| 383 | for t in self._traj: |
| 384 | self._mask.__setitem__(t, 0) |
| 385 | |
| 386 | def _update_trajectory(self, xm, ym): |
| 387 | """Update current trajectory position in mask. |
| 388 | |
| 389 | If the new position has already been filled, raise `InvalidIndexError`. |
| 390 | """ |
| 391 | if self._current_xy != (xm, ym): |
| 392 | if self[ym, xm] == 0: |
| 393 | self._traj.append((ym, xm)) |
| 394 | self._mask[ym, xm] = 1 |
| 395 | self._current_xy = (xm, ym) |
| 396 | else: |
| 397 | raise InvalidIndexError |
| 398 | |
| 399 | |
| 400 | class InvalidIndexError(Exception): |