diff --git a/python/streamtracer/streamline.py b/python/streamtracer/streamline.py index f00cbe1..fea6255 100755 --- a/python/streamtracer/streamline.py +++ b/python/streamtracer/streamline.py @@ -9,29 +9,30 @@ class VectorGrid: """ A grid of vectors. + .. note:: + + If any of *cyclic* are ``True``, then the grid values on each side of the + cyclic dimension **must** match, e.g. if ``cyclic=[False, True, False]``, + ``vectors[:, 0, :, :]`` must equal ``vectors[:, -1, :, :]``. + Parameters ---------- - vectors : array + vectors : array-like A (nx, ny, nz, 3) shaped array. The three values at (i, j, k, :) specify the (x, y, z) components of the vector at index (i, j, k). - grid_spacing : array, optional + grid_spacing : array-like, optional A (3,) shaped array, that contains the grid spacings in the (x, y, z) directions. If not specified ``grid_coords`` must be specified. - origin_coord = [float, float, float], optional + origin_coord : [`float`, `float`, `float`], optional The coordinate of the ``vectors[0, 0, 0, :]`` vector at the corner of - the box. Defaults to ``[0, 0, 0]``. - cyclic : [bool, bool, bool], optional + the box. Defaults to ``[0, 0, 0]``. This is not used if ``grid_coords`` + is specified. + cyclic : [`bool`, `bool`, `bool`], optional Whether to have cyclic boundary conditions in each of the (x, y, z) directions. Defaults to ``[False, False, False]``. grid_coords : list[array], optional - A len(3) list storing the {x, y, z} coordinates of the grid. If not + A list of length 3 storing the (x, y, z) coordinates of the grid. If not specified ``grid_spacing`` must be specified. - - Notes - ----- - If any of *cyclic* are ``True``, then the grid values on each side of the - cyclic dimension **must** match, e.g. if ``cyclic=[False, True, False]``, - ``vectors[:, 0, :, :]`` must equal ``vectors[:, -1, :, :]``. """ def __init__( @@ -45,94 +46,125 @@ def __init__( ): if grid_spacing is not None and grid_coords is not None: raise ValueError( - 'Only one of "grid_spacing" and "grid_coords" ' "can be specified." + 'Only one of "grid_spacing" and "grid_coords" can be specified.' ) if grid_spacing is None and grid_coords is None: raise ValueError( - 'One of "grid_spacing" and "grid_coords" must ' "be specified." + 'One of "grid_spacing" and "grid_coords" must be specified.' + ) + if grid_coords is not None and origin_coord is not None: + raise ValueError( + 'Specifying both "grid_coords" and "origin_coord" is ambiguous.' ) - - if cyclic is None: - cyclic = [False, False, False] - if origin_coord is None: - origin_coord = [0, 0, 0] - - if grid_spacing is not None: - grid_spacing = np.array(grid_spacing) - self._validate_spacing(grid_spacing) - elif grid_coords is not None: - self._validate_coords(grid_coords, vectors) - - self._validate_vectors(vectors) - self._validate_cyclic(vectors, cyclic) - - self.vectors = vectors self.grid_spacing = grid_spacing + self.vectors = vectors + self.cyclic = cyclic self.coords = grid_coords - self.cyclic = np.array(cyclic, dtype=bool) + self.origin_coord = origin_coord + + @property + def grid_spacing(self): + """ + Physical spacing between grid points along each axis. + """ + return self._grid_spacing + + @grid_spacing.setter + def grid_spacing(self, val): + if val is not None: + val = np.array(val) + if val.shape != (3,): + raise ValueError( + f"grid spacing must have shape (3,), got " f"{val.shape}" + ) + self._grid_spacing = val - self._origin_coord = np.array(origin_coord) + @property + def vectors(self): + """ + Three-dimensional vector field through which the streamlines will be traced. + """ + return self._vectors - @staticmethod - def _validate_vectors(vectors): - if len(vectors.shape) != 4: + @vectors.setter + def vectors(self, val): + if len(val.shape) != 4: raise ValueError("vectors must be a 4D array") - if vectors.shape[-1] != 3: + if val.shape[-1] != 3: raise ValueError( - "vectors must have shape (nx, ny, nz, 3), " f"got {vectors.shape}" + "vectors must have shape (nx, ny, nz, 3), " f"got {val.shape}" ) + self._vectors = val - @staticmethod - def _validate_spacing(grid_spacing): - if grid_spacing.shape != (3,): - raise ValueError( - f"grid spacing must have shape (3,), got " f"{grid_spacing.shape}" - ) + @property + def coords(self): + """ + The physical coordinates along each axis of the grid. + """ + return self._coords + + @coords.setter + def coords(self, val): + if val is not None: + if len(val) != 3: + raise ValueError("coords must be len(3)") + for i, dim in zip(range(3), ["x", "y", "z"]): + shape = np.array(val[i]).shape + if shape != (self.vectors.shape[i],): + raise ValueError( + f"Expected {self.vectors.shape[i]} {dim} " + f"coordinates but got {shape}" + ) + self._coords = val - @staticmethod - def _validate_coords(coords, vectors): - if len(coords) != 3: - raise ValueError("coords must be len(3)") - for i, dim in zip(range(3), ["x", "y", "z"]): - shape = np.array(coords[i]).shape - if shape != (vectors.shape[i],): - raise ValueError( - f"Expected {vectors.shape[i]} {dim} " f"coordinates but got {shape}" - ) + @property + def cyclic(self): + """ + Boolean describing whether to have cyclic boundary conditions in each of the (x, y, z) + directions. + """ + return self._cyclic - @staticmethod - def _validate_cyclic(vectors, cyclic): + @cyclic.setter + def cyclic(self, val): + if val is None: + val = [False, False, False] dims = {0: "x", 1: "y", 2: "z"} s = [slice(None)] * 4 - for i, c in enumerate(cyclic): + for i, c in enumerate(val): if c: slc = s.copy() slc[i] = slice(0, 1) - side1 = vectors[tuple(slc)] + side1 = self.vectors[tuple(slc)] slc[i] = slice(-1, None) - side2 = vectors[tuple(slc)] + side2 = self.vectors[tuple(slc)] np.testing.assert_equal( side1, side2, - err_msg=f"grid values in dimension {dims[i]} (size {vectors.shape[i]}) " + err_msg=f"grid values in dimension {dims[i]} (size {self.vectors.shape[i]}) " "do not match on each side of the cube", ) - - @property - def cyclic(self): - return self._cyclic - - @cyclic.setter - def cyclic(self, val): self._cyclic = np.array(val, dtype=bool) @property def origin_coord(self): - if self.grid_spacing is not None: - return self._origin_coord + """ + The physical coordinate corresponding to the index at ``(0,0,0)``. + """ + return self._origin_coord + + @origin_coord.setter + def origin_coord(self, val): + if val is None: + if self.grid_spacing is not None: + self._origin_coord = np.array([0, 0, 0]) + else: + self._origin_coord = np.array( + [self.xcoords[0], self.ycoords[0], self.zcoords[0]] + ) else: - return np.array([self.xcoords[0], self.ycoords[0], self.zcoords[0]]) + self._origin_coord = np.array(val) def _get_coords(self, i): if self.grid_spacing is not None: @@ -146,21 +178,21 @@ def _get_coords(self, i): @property def xcoords(self): """ - Coordinates of the x grid points. + Physical coordinates corresponding to grid points in the x-direction. """ return self._get_coords(0) @property def ycoords(self): """ - Coordinates of the x grid points. + Physical coordinates corresponding to grid points in the y-direction. """ return self._get_coords(1) @property def zcoords(self): """ - Coordinates of the x grid points. + Physical coordinates corresponding to grid points in the z-direction. """ return self._get_coords(2) @@ -171,35 +203,59 @@ class StreamTracer: Parameters ---------- - max_steps : int + max_steps : `int` Number of steps available for each line. The maximum number of points on a single stream line is ``max_steps``. - step_size : float + step_size : `float` Step size as a the fraction of cell size. - cyclic : [bool, bool, bool], optional - Whether to have cyclic boundary conditions in each dimension. + """ - Attributes - ---------- - xs : array of (n, 3) arrays - An array of the streamlines, which in general can have varying - numbers of points. - ROT : integer array - Reason(s) of termination. Shape ``len(xs)`` if traced in one direction, + def __init__(self, max_steps, step_size): + self.max_steps = max_steps + self.ds = step_size + self.xs = None + + @property + def xs(self): + """ + List of the coordinates along each streamline. + + List of three-dimensional coordinates for each streamline. + Each strealine coordinate has a shape ``(n,3)``, where ``n`` can, in principle, vary + from one streamline to the next. + """ + return self._xs + + @xs.setter + def xs(self, val): + self._xs = val + + @property + def ROT(self): + """ + Reason(s) of termination. + + Integer array with shape ``len(xs)`` if traced in one direction, or ``(len(xs), 2)`` if traced in both directions. Can take the following values: + - -1: Encountered a NaN - 1: Reached maximum available steps - 2: Out of bounds - """ + """ + return self._ROT - def __init__(self, max_steps, step_size): - self.max_steps = max_steps - self.ds = step_size - self.xs = None + @ROT.setter + def ROT(self, val): + self._ROT = val @property def max_steps(self): + """ + Number of steps available for each line. + + The maximum number of points on a single stream line is ``max_steps``. + """ return self._max_steps @max_steps.setter @@ -212,7 +268,6 @@ def max_steps(self, val): self._max_steps = val - # Calculate the streamline from a vector array def trace(self, seeds, grid, direction=0): """ Trace streamlines. @@ -222,11 +277,11 @@ def trace(self, seeds, grid, direction=0): Parameters ---------- - seeds : (n, 3) array + seeds : array-like with shape ``(n, 3)`` Seed points. - grid : VectorGrid + grid : `VectorGrid` Grid of field vectors. - direction : int, optional + direction : `int`, optional Integration direction. ``0`` for both directions, ``1`` for forward, or ``-1`` for backwards. """ @@ -268,7 +323,7 @@ def trace(self, seeds, grid, direction=0): self.xs += grid.origin_coord # Reduce the size of the arrays - self.xs = np.array([xi[:ni, :] for xi, ni in zip(self.xs, self.ns)]) + self.xs = [xi[:ni, :] for xi, ni in zip(self.xs, self.ns)] elif direction == 0: # Calculate forward streamline