jan/gridlock
2
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

move to simpler docstring format and improve type annotations

Browse Source
This commit is contained in:
Jan Petykiewicz 2021-01-08 22:43:22 -08:00
parent 4e0a5ede21
commit db9e7831b9
5 changed files with 224 additions and 170 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
gridlock/_helpers.py
View File

@ -3,8 +3,6 @@ from typing import Any
def is_scalar(var: Any) -> bool: def is_scalar(var: Any) -> bool:
""" """
Alias for 'not hasattr(var, "__len__")' Alias for `not hasattr(var, "__len__")`
:param var: Checks if var has a length.
""" """
return not hasattr(var, "__len__") return not hasattr(var, "__len__")

135
gridlock/draw.py
View File

@ -1,9 +1,9 @@
""" """
Drawing-related methods for Grid class Drawing-related methods for Grid class
""" """
from typing import List from typing import List, Optional, Union, Sequence, Callable
import numpy import numpy # type: ignore
from numpy import diff, floor, ceil, zeros, hstack, newaxis from numpy import diff, floor, ceil, zeros, hstack, newaxis
from float_raster import raster from float_raster import raster
@ -12,27 +12,34 @@ from . import GridError, Direction
from ._helpers import is_scalar from ._helpers import is_scalar
eps_callable_t = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray]
def draw_polygons(self, def draw_polygons(self,
surface_normal: Direction or int, surface_normal: Union[Direction, int],
center: List or numpy.ndarray, center: numpy.ndarray,
polygons: List[numpy.ndarray or List], polygons: Sequence[numpy.ndarray],
thickness: float, thickness: float,
eps: List[float or eps_callable_type] or float or eps_callable_type): eps: Union[Sequence[Union[float, eps_callable_t]], float, eps_callable_t],
) -> None:
""" """
Draw polygons on an axis-aligned plane. Draw polygons on an axis-aligned plane.
:param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or Args:
integer in range(3) surface_normal: Axis normal to the plane we're drawing on. Can be a `Direction` or
:param center: 3-element ndarray or list specifying an offset applied to all the polygons integer in `range(3)`
:param polygons: List of Nx2 or Nx3 ndarrays, each specifying the vertices of a polygon center: 3-element ndarray or list specifying an offset applied to all the polygons
(non-closed, clockwise). If Nx3, the surface_normal coordinate is ignored. Each polygon polygons: List of Nx2 or Nx3 ndarrays, each specifying the vertices of a polygon
must have at least 3 vertices. (non-closed, clockwise). If Nx3, the surface_normal coordinate is ignored. Each
:param thickness: Thickness of the layer to draw polygon must have at least 3 vertices.
:param eps: Value to draw with ('epsilon'). Can be scalar, callable, or a list thickness: Thickness of the layer to draw
of any of these (1 per grid). Callable values should take ndarrays x, y, z of equal eps: Value to draw with ('epsilon'). Can be scalar, callable, or a list
shape and return an ndarray of equal shape containing the eps value at the given x, y, of any of these (1 per grid). Callable values should take an ndarray the shape of the
grid and return an ndarray of equal shape containing the eps value at the given x, y,
and z (natural, not grid coordinates). and z (natural, not grid coordinates).
:raises: GridError
Raises:
GridError
""" """
# Turn surface_normal into its integer representation # Turn surface_normal into its integer representation
if isinstance(surface_normal, Direction): if isinstance(surface_normal, Direction):
@ -180,39 +187,43 @@ def draw_polygons(self,
def draw_polygon(self, def draw_polygon(self,
surface_normal: Direction or int, surface_normal: Union[Direction, int],
center: List or numpy.ndarray, center: numpy.ndarray,
polygon: List or numpy.ndarray, polygon: numpy.ndarray,
thickness: float, thickness: float,
eps: List[float or eps_callable_type] or float or eps_callable_type): eps: Union[Sequence[Union[float, eps_callable_t]], float, eps_callable_t],
) -> None:
""" """
Draw a polygon on an axis-aligned plane. Draw a polygon on an axis-aligned plane.
:param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or Args:
surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or
integer in range(3) integer in range(3)
:param center: 3-element ndarray or list specifying an offset applied to the polygon center: 3-element ndarray or list specifying an offset applied to the polygon
:param polygon: Nx2 or Nx3 ndarray specifying the vertices of a polygon (non-closed, polygon: Nx2 or Nx3 ndarray specifying the vertices of a polygon (non-closed,
clockwise). If Nx3, the surface_normal coordinate is ignored. Must have at least 3 clockwise). If Nx3, the surface_normal coordinate is ignored. Must have at
vertices. least 3 vertices.
:param thickness: Thickness of the layer to draw thickness: Thickness of the layer to draw
:param eps: Value to draw with ('epsilon'). See draw_polygons() for details. eps: Value to draw with ('epsilon'). See `draw_polygons()` for details.
""" """
self.draw_polygons(surface_normal, center, [polygon], thickness, eps) self.draw_polygons(surface_normal, center, [polygon], thickness, eps)
def draw_slab(self, def draw_slab(self,
surface_normal: Direction or int, surface_normal: Union[Direction, int],
center: List or numpy.ndarray, center: numpy.ndarray,
thickness: float, thickness: float,
eps: List[float or eps_callable_type] or float or eps_callable_type): eps: Union[List[Union[float, eps_callable_t]], float, eps_callable_t],
) -> None:
""" """
Draw an axis-aligned infinite slab. Draw an axis-aligned infinite slab.
:param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or Args:
integer in range(3) surface_normal: Axis normal to the plane we're drawing on. Can be a `Direction` or
:param center: Surface_normal coordinate at the center of the slab integer in `range(3)`
:param thickness: Thickness of the layer to draw center: Surface_normal coordinate at the center of the slab
:param eps: Value to draw with ('epsilon'). See draw_polygons() for details. thickness: Thickness of the layer to draw
eps: Value to draw with ('epsilon'). See `draw_polygons()` for details.
""" """
# Turn surface_normal into its integer representation # Turn surface_normal into its integer representation
if isinstance(surface_normal, Direction): if isinstance(surface_normal, Direction):
@ -250,16 +261,18 @@ def draw_slab(self,
def draw_cuboid(self, def draw_cuboid(self,
center: List or numpy.ndarray, center: numpy.ndarray,
dimensions: List or numpy.ndarray, dimensions: numpy.ndarray,
eps: List[float or eps_callable_type] or float or eps_callable_type): eps: Union[List[Union[float, eps_callable_t]], float, eps_callable_t],
) -> None:
""" """
Draw an axis-aligned cuboid Draw an axis-aligned cuboid
:param center: 3-element ndarray or list specifying the cuboid's center Args:
:param dimensions: 3-element list or ndarray containing the x, y, and z edge-to-edge center: 3-element ndarray or list specifying the cuboid's center
dimensions: 3-element list or ndarray containing the x, y, and z edge-to-edge
sizes of the cuboid sizes of the cuboid
:param eps: Value to draw with ('epsilon'). See draw_polygons() for details. eps: Value to draw with ('epsilon'). See `draw_polygons()` for details.
""" """
p = numpy.array([[-dimensions[0], +dimensions[1]], p = numpy.array([[-dimensions[0], +dimensions[1]],
[+dimensions[0], +dimensions[1]], [+dimensions[0], +dimensions[1]],
@ -270,22 +283,24 @@ def draw_cuboid(self,
def draw_cylinder(self, def draw_cylinder(self,
surface_normal: Direction or int, surface_normal: Union[Direction, int],
center: List or numpy.ndarray, center: numpy.ndarray,
radius: float, radius: float,
thickness: float, thickness: float,
num_points: int, num_points: int,
eps: List[float or eps_callable_type] or float or eps_callable_type): eps: Union[List[Union[float, eps_callable_t]], float, eps_callable_t],
) -> None:
""" """
Draw an axis-aligned cylinder. Approximated by a num_points-gon Draw an axis-aligned cylinder. Approximated by a num_points-gon
:param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or Args:
integer in range(3) surface_normal: Axis normal to the plane we're drawing on. Can be a `Direction` or
:param center: 3-element ndarray or list specifying the cylinder's center integer in `range(3)`
:param radius: cylinder radius center: 3-element ndarray or list specifying the cylinder's center
:param thickness: Thickness of the layer to draw radius: cylinder radius
:param num_points: The circle is approximated by a polygon with num_points vertices thickness: Thickness of the layer to draw
:param eps: Value to draw with ('epsilon'). See draw_polygons() for details. num_points: The circle is approximated by a polygon with `num_points` vertices
eps: Value to draw with ('epsilon'). See `draw_polygons()` for details.
""" """
theta = numpy.linspace(0, 2*numpy.pi, num_points, endpoint=False) theta = numpy.linspace(0, 2*numpy.pi, num_points, endpoint=False)
x = radius * numpy.sin(theta) x = radius * numpy.sin(theta)
@ -295,17 +310,19 @@ def draw_cylinder(self,
def draw_extrude_rectangle(self, def draw_extrude_rectangle(self,
rectangle: List or numpy.ndarray, rectangle: numpy.ndarray,
direction: Direction or int, direction: Union[Direction, int],
polarity: int, polarity: int,
distance: float): distance: float,
) -> None:
""" """
Extrude a rectangle of a previously-drawn structure along an axis. Extrude a rectangle of a previously-drawn structure along an axis.
:param rectangle: 2x3 ndarray or list specifying the rectangle's corners Args:
:param direction: Direction to extrude in. Direction enum or int in range(3) rectangle: 2x3 ndarray or list specifying the rectangle's corners
:param polarity: +1 or -1, direction along axis to extrude in direction: Direction to extrude in. Direction enum or int in range(3)
:param distance: How far to extrude polarity: +1 or -1, direction along axis to extrude in
distance: How far to extrude
""" """
# Turn extrude_direction into its integer representation # Turn extrude_direction into its integer representation
if isinstance(direction, Direction): if isinstance(direction, Direction):

129
gridlock/grid.py
View File

@ -1,6 +1,6 @@
from typing import List, Tuple, Callable, Dict from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar
import numpy import numpy # type: ignore
from numpy import diff, floor, ceil, zeros, hstack, newaxis from numpy import diff, floor, ceil, zeros, hstack, newaxis
import pickle import pickle
@ -23,24 +23,26 @@ class Grid(object):
is generated based on the coordinates of the boundary points). Also does is generated based on the coordinates of the boundary points). Also does
straightforward natural <-> grid unit conversion. straightforward natural <-> grid unit conversion.
self.grids[i][a,b,c] contains the value of epsilon for the cell located around `self.grids[i][a,b,c]` contains the value of epsilon for the cell located around
```
(xyz[0][a] + dxyz[0][a] * shifts[i, 0], (xyz[0][a] + dxyz[0][a] * shifts[i, 0],
xyz[1][b] + dxyz[1][b] * shifts[i, 1], xyz[1][b] + dxyz[1][b] * shifts[i, 1],
xyz[2][c] + dxyz[2][c] * shifts[i, 2]). xyz[2][c] + dxyz[2][c] * shifts[i, 2]).
You can get raw edge coordinates (exyz), ```
center coordinates (xyz), You can get raw edge coordinates (`exyz`),
cell sizes (dxyz), center coordinates (`xyz`),
cell sizes (`dxyz`),
from the properties named as above, or get them for a given grid by using the from the properties named as above, or get them for a given grid by using the
self.shifted_*xyz(which_shifts) functions. `self.shifted_*xyz(which_shifts)` functions.
The sizes of adjacent cells are taken into account when applying shifts. The The sizes of adjacent cells are taken into account when applying shifts. The
total shift for each edge is chosen using (shift * dx_of_cell_being_moved_through). total shift for each edge is chosen using `(shift * dx_of_cell_being_moved_through)`.
It is tricky to determine the size of the right-most cell after shifting, It is tricky to determine the size of the right-most cell after shifting,
since its right boundary should shift by shifts[i][a] * dxyz[a][dxyz[a].size], since its right boundary should shift by `shifts[i][a] * dxyz[a][dxyz[a].size]`,
where the dxyz element refers to a cell that does not exist. where the dxyz element refers to a cell that does not exist.
Because of this, we either assume this 'ghost' cell is the same size as the last Because of this, we either assume this 'ghost' cell is the same size as the last
real cell, or, if self.periodic[a] is set to True, the same size as the first cell. real cell, or, if `self.periodic[a]` is set to `True`, the same size as the first cell.
""" """
Yee_Shifts_E = 0.5 * numpy.array([[1, 0, 0], Yee_Shifts_E = 0.5 * numpy.array([[1, 0, 0],
@ -63,7 +65,8 @@ class Grid(object):
""" """
Cell sizes for each axis, no shifts applied Cell sizes for each axis, no shifts applied
:return: List of 3 ndarrays of cell sizes Returns:
List of 3 ndarrays of cell sizes
""" """
return [diff(self.exyz[a]) for a in range(3)] return [diff(self.exyz[a]) for a in range(3)]
@ -72,7 +75,8 @@ class Grid(object):
""" """
Cell centers for each axis, no shifts applied Cell centers for each axis, no shifts applied
:return: List of 3 ndarrays of cell edges Returns:
List of 3 ndarrays of cell edges
""" """
return [self.exyz[a][:-1] + self.dxyz[a] / 2.0 for a in range(3)] return [self.exyz[a][:-1] + self.dxyz[a] / 2.0 for a in range(3)]
@ -81,7 +85,8 @@ class Grid(object):
""" """
The number of cells in x, y, and z The number of cells in x, y, and z
:return: ndarray [x_centers.size, y_centers.size, z_centers.size] Returns:
ndarray of [x_centers.size, y_centers.size, z_centers.size]
""" """
return numpy.array([coord.size - 1 for coord in self.exyz], dtype=int) return numpy.array([coord.size - 1 for coord in self.exyz], dtype=int)
@ -95,7 +100,8 @@ class Grid(object):
If periodic, final edge shifts same amount as first If periodic, final edge shifts same amount as first
Otherwise, final edge shifts same amount as second-to-last Otherwise, final edge shifts same amount as second-to-last
:return: list of [dxs, dys, dzs] with each element same length as elements of self.xyz Returns:
list of [dxs, dys, dzs] with each element same length as elements of `self.xyz`
""" """
el = [0 if p else -1 for p in self.periodic] el = [0 if p else -1 for p in self.periodic]
return [hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)] return [hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)]
@ -104,7 +110,9 @@ class Grid(object):
def center(self) -> numpy.ndarray: def center(self) -> numpy.ndarray:
""" """
Center position of the entire grid, no shifts applied Center position of the entire grid, no shifts applied
:return: ndarray [x_center, y_center, z_center]
Returns:
ndarray of [x_center, y_center, z_center]
""" """
# center is just average of first and last xyz, which is just the average of the # center is just average of first and last xyz, which is just the average of the
# first two and last two exyz # first two and last two exyz
@ -118,18 +126,22 @@ class Grid(object):
ndarrays. No shifts are applied, so these are extreme bounds on these values (as a ndarrays. No shifts are applied, so these are extreme bounds on these values (as a
weighted average is performed when shifting). weighted average is performed when shifting).
:return: List of 2 ndarrays, d_min=[min(dx), min(dy), min(dz)] and d_max=[...] Returns:
Tuple of 2 ndarrays, `d_min=[min(dx), min(dy), min(dz)]` and `d_max=[...]`
""" """
d_min = numpy.array([min(self.dxyz[a]) for a in range(3)], dtype=float) d_min = numpy.array([min(self.dxyz[a]) for a in range(3)], dtype=float)
d_max = numpy.array([max(self.dxyz[a]) for a in range(3)], dtype=float) d_max = numpy.array([max(self.dxyz[a]) for a in range(3)], dtype=float)
return d_min, d_max return d_min, d_max
def shifted_exyz(self, which_shifts: int or None) -> List[numpy.ndarray]: def shifted_exyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]:
""" """
Returns edges for which_shifts. Returns edges for which_shifts.
:param which_shifts: Which grid (which shifts) to use, or None for unshifted Args:
:return: List of 3 ndarrays of cell edges which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell edges
""" """
if which_shifts is None: if which_shifts is None:
return self.exyz return self.exyz
@ -143,12 +155,15 @@ class Grid(object):
return [self.exyz[a] + dxyz[a] * shifts[a] for a in range(3)] return [self.exyz[a] + dxyz[a] * shifts[a] for a in range(3)]
def shifted_dxyz(self, which_shifts: int or None) -> List[numpy.ndarray]: def shifted_dxyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]:
""" """
Returns cell sizes for which_shifts. Returns cell sizes for `which_shifts`.
:param which_shifts: Which grid (which shifts) to use, or None for unshifted Args:
:return: List of 3 ndarrays of cell sizes which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell sizes
""" """
if which_shifts is None: if which_shifts is None:
return self.dxyz return self.dxyz
@ -167,12 +182,15 @@ class Grid(object):
return sdxyz return sdxyz
def shifted_xyz(self, which_shifts: int or None) -> List[numpy.ndarray]: def shifted_xyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]:
""" """
Returns cell centers for which_shifts. Returns cell centers for `which_shifts`.
:param which_shifts: Which grid (which shifts) to use, or None for unshifted Args:
:return: List of 3 ndarrays of cell centers which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell centers
""" """
if which_shifts is None: if which_shifts is None:
return self.xyz return self.xyz
@ -184,7 +202,8 @@ class Grid(object):
""" """
Return cell widths, with each dimension shifted by the corresponding shifts. Return cell widths, with each dimension shifted by the corresponding shifts.
:return: [grid.shifted_dxyz(which_shifts=a)[a] for a in range(3)] Returns:
`[grid.shifted_dxyz(which_shifts=a)[a] for a in range(3)]`
""" """
if len(self.grids) != 3: if len(self.grids) != 3:
raise GridError('autoshifting requires exactly 3 grids') raise GridError('autoshifting requires exactly 3 grids')
@ -192,29 +211,31 @@ class Grid(object):
def __init__(self, def __init__(self,
pixel_edge_coordinates: List[List or numpy.ndarray], pixel_edge_coordinates: Sequence[numpy.ndarray],
shifts: numpy.ndarray or List = Yee_Shifts_E, shifts: numpy.ndarray = Yee_Shifts_E,
initial: float or numpy.ndarray or List[float] or List[numpy.ndarray] = 1.0, initial: Union[float, numpy.ndarray] = 1.0,
num_grids: int = None, num_grids: Optional[int] = None,
periodic: bool or List[bool] = False): periodic: Union[bool, Sequence[bool]] = False,
) -> None:
""" """
Initialize a new Grid Args:
pixel_edge_coordinates: 3-element list of (ndarrays or lists) specifying the
:param pixel_edge_coordinates: 3-element list of (ndarrays or lists) specifying the
coordinates of the pixel edges in each dimensions coordinates of the pixel edges in each dimensions
(ie, [[x0, x1, x2,...], [y0,...], [z0,...]] where the first pixel has x-edges x=x0 and (ie, `[[x0, x1, x2,...], [y0,...], [z0,...]]` where the first pixel has x-edges x=`x0` and
x=x1, the second has edges x=x1 and x=x2, etc.) x=`x1`, the second has edges x=`x1` and x=`x2`, etc.)
:param shifts: Nx3 array containing [x, y, z] offsets for each of N grids. shifts: Nx3 array containing `[x, y, z]` offsets for each of N grids.
E-field Yee shifts are used by default. E-field Yee shifts are used by default.
:param initial: Grids are initialized to this value. If scalar, all grids are initialized initial: Grids are initialized to this value. If scalar, all grids are initialized
with ndarrays full of the scalar. If a list of scalars, grid[i] is initialized to an with ndarrays full of the scalar. If a list of scalars, `grid[i]` is initialized to an
ndarray full of initial[i]. If a list of ndarrays of the same shape as the grids, grid[i] ndarray full of `initial[i]`. If a list of ndarrays of the same shape as the grids, `grid[i]`
is set to initial[i]. Default 1. is set to `initial[i]`. Default `1.0`.
:param num_grids: How many grids to create. Must be <= shifts.shape[0]. num_grids: How many grids to create. Must be <= `shifts.shape[0]`.
Default is shifts.shape[0] Default is `shifts.shape[0]`
:param periodic: Specifies how the sizes of edge cells are calculated; see main class periodic: Specifies how the sizes of edge cells are calculated; see main class
documentation. List of 3 bool, or a single bool that gets broadcast. Default False. documentation. List of 3 bool, or a single bool that gets broadcast. Default `False`.
:raises: GridError
Raises:
`GridError` on invalid input
""" """
self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)] # type: List[numpy.ndarray] self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)] # type: List[numpy.ndarray]
"""Cell edges. Monotonically increasing without duplicates.""" """Cell edges. Monotonically increasing without duplicates."""
@ -280,7 +301,8 @@ class Grid(object):
""" """
Load a grid from a file Load a grid from a file
:param filename: Filename to load from. Args:
filename: Filename to load from.
""" """
with open(filename, 'rb') as f: with open(filename, 'rb') as f:
tmp_dict = pickle.load(f) tmp_dict = pickle.load(f)
@ -293,16 +315,15 @@ class Grid(object):
""" """
Save to file. Save to file.
:param filename: Filename to save to. Args:
filename: Filename to save to.
""" """
with open(filename, 'wb') as f: with open(filename, 'wb') as f:
pickle.dump(self.__dict__, f, protocol=2) pickle.dump(self.__dict__, f, protocol=2)
def copy(self): def copy(self):
""" """
Return a deep copy of the grid. Returns:
Deep copy of the grid.
:return: Deep copy of the grid.
""" """
return copy.deepcopy(self) return copy.deepcopy(self)

59
gridlock/position.py
View File

@ -1,34 +1,39 @@
""" """
Position-related methods for Grid class Position-related methods for Grid class
""" """
from typing import List, Optional
from typing import List import numpy # type: ignore
import numpy
from numpy import zeros from numpy import zeros
from . import GridError from . import GridError
def ind2pos(self, def ind2pos(self,
ind: numpy.ndarray or List, ind: numpy.ndarray,
which_shifts: int = None, which_shifts: Optional[int] = None,
round_ind: bool = True, round_ind: bool = True,
check_bounds: bool = True check_bounds: bool = True
) -> numpy.ndarray: ) -> numpy.ndarray:
""" """
Returns the natural position corresponding to the specified cell center indices. Returns the natural position corresponding to the specified cell center indices.
The resulting position is clipped to the bounds of the grid The resulting position is clipped to the bounds of the grid
(to cell centers if round_ind=True, or cell outer edges if round_ind=False) (to cell centers if `round_ind=True`, or cell outer edges if `round_ind=False`)
:param ind: Indices of the position. Can be fractional. (3-element ndarray or list) Args:
:param which_shifts: which grid number (shifts) to use ind: Indices of the position. Can be fractional. (3-element ndarray or list)
:param round_ind: Whether to round ind to the nearest integer position before indexing which_shifts: which grid number (`shifts`) to use
(default True) round_ind: Whether to round ind to the nearest integer position before indexing
:param check_bounds: Whether to raise an GridError if the provided ind is outside of (default `True`)
the grid, as defined above (centers if round_ind, else edges) (default True) check_bounds: Whether to raise an `GridError` if the provided ind is outside of
:return: 3-element ndarray specifying the natural position the grid, as defined above (centers if `round_ind`, else edges) (default `True`)
:raises: GridError
Returns:
3-element ndarray specifying the natural position
Raises:
`GridError` if invalid `which_shifts`
`GridError` if `check_bounds` and out of bounds
""" """
if which_shifts is not None and which_shifts >= self.shifts.shape[0]: if which_shifts is not None and which_shifts >= self.shifts.shape[0]:
raise GridError('Invalid shifts') raise GridError('Invalid shifts')
@ -56,21 +61,27 @@ def ind2pos(self,
def pos2ind(self, def pos2ind(self,
r: numpy.ndarray or List, r: numpy.ndarray,
which_shifts: int or None, which_shifts: Optional[int],
round_ind: bool=True, round_ind: bool = True,
check_bounds: bool=True check_bounds: bool = True
) -> numpy.ndarray: ) -> numpy.ndarray:
""" """
Returns the cell-center indices corresponding to the specified natural position. Returns the cell-center indices corresponding to the specified natural position.
The resulting position is clipped to within the outer centers of the grid. The resulting position is clipped to within the outer centers of the grid.
:param r: Natural position that we will convert into indices (3-element ndarray or list) Args:
:param which_shifts: which grid number (shifts) to use r: Natural position that we will convert into indices (3-element ndarray or list)
:param round_ind: Whether to round the returned indices to the nearest integers. which_shifts: which grid number (`shifts`) to use
:param check_bounds: Whether to throw an GridError if r is outside the grid edges round_ind: Whether to round the returned indices to the nearest integers.
:return: 3-element ndarray specifying the indices check_bounds: Whether to throw an `GridError` if `r` is outside the grid edges
:raises: GridError
Returns:
3-element ndarray specifying the indices
Raises:
`GridError` if invalid `which_shifts`
`GridError` if `check_bounds` and out of bounds
""" """
r = numpy.squeeze(r) r = numpy.squeeze(r)
if r.size != 3: if r.size != 3:

55
gridlock/read.py
View File

@ -1,9 +1,9 @@
""" """
Readback and visualization methods for Grid class Readback and visualization methods for Grid class
""" """
from typing import Dict from typing import Dict, Optional, Union, Any
import numpy import numpy # type: ignore
from numpy import floor, ceil, zeros from numpy import floor, ceil, zeros
from . import GridError, Direction from . import GridError, Direction
@ -15,7 +15,7 @@ from ._helpers import is_scalar
def get_slice(self, def get_slice(self,
surface_normal: Direction or int, surface_normal: Union[Direction, int],
center: float, center: float,
which_shifts: int = 0, which_shifts: int = 0,
sample_period: int = 1 sample_period: int = 1
@ -24,12 +24,15 @@ def get_slice(self,
Retrieve a slice of a grid. Retrieve a slice of a grid.
Interpolates if given a position between two planes. Interpolates if given a position between two planes.
:param surface_normal: Axis normal to the plane we're displaying. Can be a Direction or Args:
integer in range(3) surface_normal: Axis normal to the plane we're displaying. Can be a `Direction` or
:param center: Scalar specifying position along surface_normal axis. integer in `range(3)`
:param which_shifts: Which grid to display. Default is the first grid (0). center: Scalar specifying position along surface_normal axis.
:param sample_period: Period for down-sampling the image. Default 1 (disabled) which_shifts: Which grid to display. Default is the first grid (0).
:return Array containing the portion of the grid. sample_period: Period for down-sampling the image. Default 1 (disabled)
Returns:
Array containing the portion of the grid.
""" """
if not is_scalar(center) and numpy.isreal(center): if not is_scalar(center) and numpy.isreal(center):
raise GridError('center must be a real scalar') raise GridError('center must be a real scalar')
@ -77,22 +80,24 @@ def get_slice(self,
def visualize_slice(self, def visualize_slice(self,
surface_normal: Direction or int, surface_normal: Union[Direction, int],
center: float, center: float,
which_shifts: int = 0, which_shifts: int = 0,
sample_period: int = 1, sample_period: int = 1,
finalize: bool = True, finalize: bool = True,
pcolormesh_args: Dict = None): pcolormesh_args: Optional[Dict[str, Any]] = None,
) -> None:
""" """
Visualize a slice of a grid. Visualize a slice of a grid.
Interpolates if given a position between two planes. Interpolates if given a position between two planes.
:param surface_normal: Axis normal to the plane we're displaying. Can be a Direction or Args:
integer in range(3) surface_normal: Axis normal to the plane we're displaying. Can be a `Direction` or
:param center: Scalar specifying position along surface_normal axis. integer in `range(3)`
:param which_shifts: Which grid to display. Default is the first grid (0). center: Scalar specifying position along surface_normal axis.
:param sample_period: Period for down-sampling the image. Default 1 (disabled) which_shifts: Which grid to display. Default is the first grid (0).
:param finalize: Whether to call pyplot.show() after constructing the plot. Default True sample_period: Period for down-sampling the image. Default 1 (disabled)
finalize: Whether to call `pyplot.show()` after constructing the plot. Default `True`
""" """
from matplotlib import pyplot from matplotlib import pyplot
@ -125,19 +130,21 @@ def visualize_slice(self,
def visualize_isosurface(self, def visualize_isosurface(self,
level: float = None, level: Optional[float] = None,
which_shifts: int = 0, which_shifts: int = 0,
sample_period: int = 1, sample_period: int = 1,
show_edges: bool = True, show_edges: bool = True,
finalize: bool = True): finalize: bool = True,
) -> None:
""" """
Draw an isosurface plot of the device. Draw an isosurface plot of the device.
:param level: Value at which to find isosurface. Default (None) uses mean value in grid. Args:
:param which_shifts: Which grid to display. Default is the first grid (0). level: Value at which to find isosurface. Default (None) uses mean value in grid.
:param sample_period: Period for down-sampling the image. Default 1 (disabled) which_shifts: Which grid to display. Default is the first grid (0).
:param show_edges: Whether to draw triangle edges. Default True sample_period: Period for down-sampling the image. Default 1 (disabled)
:param finalize: Whether to call pyplot.show() after constructing the plot. Default True show_edges: Whether to draw triangle edges. Default `True`
finalize: Whether to call `pyplot.show()` after constructing the plot. Default `True`
""" """
from matplotlib import pyplot from matplotlib import pyplot
import skimage.measure import skimage.measure