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

typing and formatting updates

Browse Source
master
Jan Petykiewicz 2 years ago
parent 7d3b2272bc
commit ec5c77e018
6 changed files with 142 additions and 120 deletions
Show Stats Download Patch File Download Diff File

113
gridlock/draw.py
Unescape Escape View File

@ -3,7 +3,8 @@ Drawing-related methods for Grid class
""" """
from typing import List, Optional, Union, Sequence, Callable from typing import List, Optional, Union, Sequence, Callable
import numpy # type: ignore import numpy
from numpy.typing import NDArray, ArrayLike
from float_raster import raster from float_raster import raster
from . import GridError from . import GridError
@ -15,17 +16,19 @@ from . import GridError
# without having to pass `cell_data` again each time? # without having to pass `cell_data` again each time?
foreground_callable_t = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray] foreground_callable_t = Callable[[NDArray, NDArray, NDArray], NDArray]
foreground_t = Union[float, foreground_callable_t]
def draw_polygons(self, def draw_polygons(
cell_data: numpy.ndarray, self,
surface_normal: int, cell_data: NDArray,
center: numpy.ndarray, surface_normal: int,
polygons: Sequence[numpy.ndarray], center: ArrayLike,
thickness: float, polygons: Sequence[NDArray],
foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t], thickness: float,
) -> None: foreground: Union[Sequence[foreground_t], foreground_t],
) -> None:
""" """
Draw polygons on an axis-aligned plane. Draw polygons on an axis-aligned plane.
@ -74,8 +77,8 @@ def draw_polygons(self,
# ## Compute sub-domain of the grid occupied by polygons # ## Compute sub-domain of the grid occupied by polygons
# 1) Compute outer bounds (bd) of polygons # 1) Compute outer bounds (bd) of polygons
bd_2d_min = [0, 0] bd_2d_min = numpy.array([0, 0])
bd_2d_max = [0, 0] bd_2d_max = numpy.array([0, 0])
for polygon in polygons: for polygon in polygons:
bd_2d_min = numpy.minimum(bd_2d_min, polygon.min(axis=0)) bd_2d_min = numpy.minimum(bd_2d_min, polygon.min(axis=0))
bd_2d_max = numpy.maximum(bd_2d_max, polygon.max(axis=0)) bd_2d_max = numpy.maximum(bd_2d_max, polygon.max(axis=0))
@ -97,7 +100,7 @@ def draw_polygons(self,
polygons = [poly + center[surface] for poly in polygons] polygons = [poly + center[surface] for poly in polygons]
# ## Generate weighing function # ## Generate weighing function
def to_3d(vector: numpy.ndarray, val: float = 0.0) -> numpy.ndarray: def to_3d(vector: NDArray, val: float = 0.0) -> NDArray[numpy.float64]:
v_2d = numpy.array(vector, dtype=float) v_2d = numpy.array(vector, dtype=float)
return numpy.insert(v_2d, surface_normal, (val,)) return numpy.insert(v_2d, surface_normal, (val,))
@ -188,14 +191,15 @@ def draw_polygons(self,
cell_data[g_slice] = (1 - w) * cell_data[g_slice] + w * foreground_i cell_data[g_slice] = (1 - w) * cell_data[g_slice] + w * foreground_i
def draw_polygon(self, def draw_polygon(
cell_data: numpy.ndarray, self,
surface_normal: int, cell_data: NDArray,
center: numpy.ndarray, surface_normal: int,
polygon: numpy.ndarray, center: ArrayLike,
thickness: float, polygon: ArrayLike,
foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t], thickness: float,
) -> None: foreground: Union[Sequence[foreground_t], foreground_t],
) -> None:
""" """
Draw a polygon on an axis-aligned plane. Draw a polygon on an axis-aligned plane.
@ -212,13 +216,14 @@ def draw_polygon(self,
self.draw_polygons(cell_data, surface_normal, center, [polygon], thickness, foreground) self.draw_polygons(cell_data, surface_normal, center, [polygon], thickness, foreground)
def draw_slab(self, def draw_slab(
cell_data: numpy.ndarray, self,
surface_normal: int, cell_data: NDArray,
center: numpy.ndarray, surface_normal: int,
thickness: float, center: ArrayLike,
foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t], thickness: float,
) -> None: foreground: Union[Sequence[foreground_t], foreground_t],
) -> None:
""" """
Draw an axis-aligned infinite slab. Draw an axis-aligned infinite slab.
@ -262,12 +267,13 @@ def draw_slab(self,
self.draw_polygon(cell_data, surface_normal, center_shift, p, thickness, foreground) self.draw_polygon(cell_data, surface_normal, center_shift, p, thickness, foreground)
def draw_cuboid(self, def draw_cuboid(
cell_data: numpy.ndarray, self,
center: numpy.ndarray, cell_data: NDArray,
dimensions: numpy.ndarray, center: ArrayLike,
foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t], dimensions: ArrayLike,
) -> None: foreground: Union[Sequence[foreground_t], foreground_t],
) -> None:
""" """
Draw an axis-aligned cuboid Draw an axis-aligned cuboid
@ -278,6 +284,7 @@ def draw_cuboid(self,
sizes of the cuboid sizes of the cuboid
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details. foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
""" """
dimensions = numpy.array(dimensions, copy=False)
p = numpy.array([[-dimensions[0], +dimensions[1]], p = numpy.array([[-dimensions[0], +dimensions[1]],
[+dimensions[0], +dimensions[1]], [+dimensions[0], +dimensions[1]],
[+dimensions[0], -dimensions[1]], [+dimensions[0], -dimensions[1]],
@ -286,15 +293,16 @@ def draw_cuboid(self,
self.draw_polygon(cell_data, 2, center, p, thickness, foreground) self.draw_polygon(cell_data, 2, center, p, thickness, foreground)
def draw_cylinder(self, def draw_cylinder(
cell_data: numpy.ndarray, self,
surface_normal: int, cell_data: NDArray,
center: numpy.ndarray, surface_normal: int,
radius: float, center: ArrayLike,
thickness: float, radius: float,
num_points: int, thickness: float,
foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t], num_points: int,
) -> None: foreground: Union[Sequence[foreground_t], foreground_t],
) -> None:
""" """
Draw an axis-aligned cylinder. Approximated by a num_points-gon Draw an axis-aligned cylinder. Approximated by a num_points-gon
@ -314,13 +322,14 @@ def draw_cylinder(self,
self.draw_polygon(cell_data, surface_normal, center, polygon, thickness, foreground) self.draw_polygon(cell_data, surface_normal, center, polygon, thickness, foreground)
def draw_extrude_rectangle(self, def draw_extrude_rectangle(
cell_data: numpy.ndarray, self,
rectangle: numpy.ndarray, cell_data: NDArray,
direction: int, rectangle: ArrayLike,
polarity: int, direction: int,
distance: float, polarity: int,
) -> None: 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.
@ -361,10 +370,10 @@ def draw_extrude_rectangle(self,
mult = [1-fpart, fpart][::s] # reverses if s negative mult = [1-fpart, fpart][::s] # reverses if s negative
foreground = mult[0] * grid[tuple(ind)] foreground = mult[0] * grid[tuple(ind)]
ind[direction] += 1 ind[direction] += 1 # type: ignore #(known safe)
foreground += mult[1] * grid[tuple(ind)] foreground += mult[1] * grid[tuple(ind)]
def f_foreground(xs, ys, zs, i=i, foreground=foreground) -> numpy.ndarray: def f_foreground(xs, ys, zs, i=i, foreground=foreground) -> NDArray[numpy.int_]:
# transform from natural position to index # transform from natural position to index
xyzi = numpy.array([self.pos2ind(qrs, which_shifts=i) xyzi = numpy.array([self.pos2ind(qrs, which_shifts=i)
for qrs in zip(xs.flat, ys.flat, zs.flat)], dtype=int) for qrs in zip(xs.flat, ys.flat, zs.flat)], dtype=int)

2
gridlock/examples/ex0.py
Unescape Escape View File

@ -1,4 +1,4 @@
import numpy # type: ignore import numpy
from gridlock import Grid from gridlock import Grid

58
gridlock/grid.py
Unescape Escape View File

@ -1,6 +1,7 @@
from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar, TypeVar from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar, TypeVar
import numpy # type: ignore import numpy
from numpy.typing import NDArray, ArrayLike
from numpy import diff, floor, ceil, zeros, hstack, newaxis from numpy import diff, floor, ceil, zeros, hstack, newaxis
import pickle import pickle
@ -10,7 +11,7 @@ import copy
from . import GridError from . import GridError
foreground_callable_type = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray] foreground_callable_type = Callable[[NDArray, NDArray, NDArray], NDArray]
T = TypeVar('T', bound='Grid') T = TypeVar('T', bound='Grid')
@ -48,23 +49,27 @@ class Grid:
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.
""" """
exyz: List[numpy.ndarray] exyz: List[NDArray]
"""Cell edges. Monotonically increasing without duplicates.""" """Cell edges. Monotonically increasing without duplicates."""
periodic: List[bool] periodic: List[bool]
"""For each axis, determines how far the rightmost boundary gets shifted. """ """For each axis, determines how far the rightmost boundary gets shifted. """
shifts: numpy.ndarray shifts: NDArray
"""Offsets `[[x0, y0, z0], [x1, y1, z1], ...]` for grid `0,1,...`""" """Offsets `[[x0, y0, z0], [x1, y1, z1], ...]` for grid `0,1,...`"""
Yee_Shifts_E: ClassVar[numpy.ndarray] = 0.5 * numpy.array([[1, 0, 0], Yee_Shifts_E: ClassVar[NDArray] = 0.5 * numpy.array([
[0, 1, 0], [1, 0, 0],
[0, 0, 1]], dtype=float) [0, 1, 0],
[0, 0, 1],
], dtype=float)
"""Default shifts for Yee grid E-field""" """Default shifts for Yee grid E-field"""
Yee_Shifts_H: ClassVar[numpy.ndarray] = 0.5 * numpy.array([[0, 1, 1], Yee_Shifts_H: ClassVar[NDArray] = 0.5 * numpy.array([
[1, 0, 1], [0, 1, 1],
[1, 1, 0]], dtype=float) [1, 0, 1],
[1, 1, 0],
], dtype=float)
"""Default shifts for Yee grid H-field""" """Default shifts for Yee grid H-field"""
from .draw import ( from .draw import (
@ -75,7 +80,7 @@ class Grid:
from .position import ind2pos, pos2ind from .position import ind2pos, pos2ind
@property @property
def dxyz(self) -> List[numpy.ndarray]: def dxyz(self) -> List[NDArray]:
""" """
Cell sizes for each axis, no shifts applied Cell sizes for each axis, no shifts applied
@ -85,7 +90,7 @@ class Grid:
return [numpy.diff(ee) for ee in self.exyz] return [numpy.diff(ee) for ee in self.exyz]
@property @property
def xyz(self) -> List[numpy.ndarray]: def xyz(self) -> List[NDArray]:
""" """
Cell centers for each axis, no shifts applied Cell centers for each axis, no shifts applied
@ -95,7 +100,7 @@ class Grid:
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)]
@property @property
def shape(self) -> numpy.ndarray: def shape(self) -> NDArray[numpy.int_]:
""" """
The number of cells in x, y, and z The number of cells in x, y, and z
@ -119,7 +124,7 @@ class Grid:
return numpy.hstack((self.num_grids, self.shape)) return numpy.hstack((self.num_grids, self.shape))
@property @property
def dxyz_with_ghost(self) -> List[numpy.ndarray]: def dxyz_with_ghost(self) -> List[NDArray]:
""" """
Gives dxyz with an additional 'ghost' cell at the end, whose value depends Gives dxyz with an additional 'ghost' cell at the end, whose value depends
on whether or not the axis has periodic boundary conditions. See main description on whether or not the axis has periodic boundary conditions. See main description
@ -135,7 +140,7 @@ class Grid:
return [numpy.hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)] return [numpy.hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)]
@property @property
def center(self) -> numpy.ndarray: def center(self) -> NDArray[numpy.float64]:
""" """
Center position of the entire grid, no shifts applied Center position of the entire grid, no shifts applied
@ -148,7 +153,7 @@ class Grid:
return numpy.array(centers, dtype=float) return numpy.array(centers, dtype=float)
@property @property
def dxyz_limits(self) -> Tuple[numpy.ndarray, numpy.ndarray]: def dxyz_limits(self) -> Tuple[NDArray, NDArray]:
""" """
Returns the minimum and maximum cell size for each axis, as a tuple of two 3-element Returns the minimum and maximum cell size for each axis, as a tuple of two 3-element
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
@ -161,7 +166,7 @@ class Grid:
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: Optional[int]) -> List[numpy.ndarray]: def shifted_exyz(self, which_shifts: Optional[int]) -> List[NDArray]:
""" """
Returns edges for which_shifts. Returns edges for which_shifts.
@ -183,7 +188,7 @@ class Grid:
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: Optional[int]) -> List[numpy.ndarray]: def shifted_dxyz(self, which_shifts: Optional[int]) -> List[NDArray]:
""" """
Returns cell sizes for `which_shifts`. Returns cell sizes for `which_shifts`.
@ -210,7 +215,7 @@ class Grid:
return sdxyz return sdxyz
def shifted_xyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]: def shifted_xyz(self, which_shifts: Optional[int]) -> List[NDArray[numpy.float64]]:
""" """
Returns cell centers for `which_shifts`. Returns cell centers for `which_shifts`.
@ -226,7 +231,7 @@ class Grid:
dxyz = self.shifted_dxyz(which_shifts) dxyz = self.shifted_dxyz(which_shifts)
return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)] return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)]
def autoshifted_dxyz(self) -> List[numpy.ndarray]: def autoshifted_dxyz(self) -> List[NDArray[numpy.float64]]:
""" """
Return cell widths, with each dimension shifted by the corresponding shifts. Return cell widths, with each dimension shifted by the corresponding shifts.
@ -237,7 +242,7 @@ class Grid:
raise GridError('Autoshifting requires exactly 3 grids') raise GridError('Autoshifting requires exactly 3 grids')
return [self.shifted_dxyz(which_shifts=a)[a] for a in range(3)] return [self.shifted_dxyz(which_shifts=a)[a] for a in range(3)]
def allocate(self, fill_value: Optional[float] = 1.0, dtype=numpy.float32) -> numpy.ndarray: def allocate(self, fill_value: Optional[float] = 1.0, dtype=numpy.float32) -> NDArray:
""" """
Allocate an ndarray for storing grid data. Allocate an ndarray for storing grid data.
@ -254,11 +259,12 @@ class Grid:
else: else:
return numpy.full(self.cell_data_shape, fill_value, dtype=dtype) return numpy.full(self.cell_data_shape, fill_value, dtype=dtype)
def __init__(self, def __init__(
pixel_edge_coordinates: Sequence[numpy.ndarray], self,
shifts: numpy.ndarray = Yee_Shifts_E, pixel_edge_coordinates: Sequence[ArrayLike],
periodic: Union[bool, Sequence[bool]] = False, shifts: ArrayLike = Yee_Shifts_E,
) -> None: periodic: Union[bool, Sequence[bool]] = False,
) -> None:
""" """
Args: Args:
pixel_edge_coordinates: 3-element list of (ndarrays or lists) specifying the pixel_edge_coordinates: 3-element list of (ndarrays or lists) specifying the

31
gridlock/position.py
Unescape Escape View File

@ -1,19 +1,21 @@
""" """
Position-related methods for Grid class Position-related methods for Grid class
""" """
from typing import List, Optional from typing import List, Optional, Sequence
import numpy # type: ignore import numpy
from numpy.typing import NDArray, ArrayLike
from . import GridError from . import GridError
def ind2pos(self, def ind2pos(
ind: numpy.ndarray, self,
which_shifts: Optional[int] = None, ind: NDArray,
round_ind: bool = True, which_shifts: Optional[int] = None,
check_bounds: bool = True round_ind: bool = True,
) -> numpy.ndarray: check_bounds: bool = True
) -> NDArray[numpy.float64]:
""" """
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
@ -59,12 +61,13 @@ def ind2pos(self,
return numpy.array(position, dtype=float) return numpy.array(position, dtype=float)
def pos2ind(self, def pos2ind(
r: numpy.ndarray, self,
which_shifts: Optional[int], r: ArrayLike,
round_ind: bool = True, which_shifts: Optional[int],
check_bounds: bool = True round_ind: bool = True,
) -> numpy.ndarray: check_bounds: bool = True
) -> NDArray[numpy.float64]:
""" """
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.

54
gridlock/read.py
Unescape Escape View File

@ -3,7 +3,8 @@ Readback and visualization methods for Grid class
""" """
from typing import Dict, Optional, Union, Any from typing import Dict, Optional, Union, Any
import numpy # type: ignore import numpy
from numpy.typing import NDArray, ArrayLike
from . import GridError from . import GridError
@ -12,13 +13,14 @@ from . import GridError
# .visualize_isosurface uses mpl_toolkits.mplot3d # .visualize_isosurface uses mpl_toolkits.mplot3d
def get_slice(self, def get_slice(
cell_data: numpy.ndarray, self,
surface_normal: int, cell_data: NDArray,
center: float, surface_normal: int,
which_shifts: int = 0, center: float,
sample_period: int = 1 which_shifts: int = 0,
) -> numpy.ndarray: sample_period: int = 1
) -> NDArray:
""" """
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.
@ -75,15 +77,16 @@ def get_slice(self,
return sliced_grid return sliced_grid
def visualize_slice(self, def visualize_slice(
cell_data: numpy.ndarray, self,
surface_normal: int, cell_data: NDArray,
center: float, surface_normal: int,
which_shifts: int = 0, center: float,
sample_period: int = 1, which_shifts: int = 0,
finalize: bool = True, sample_period: int = 1,
pcolormesh_args: Optional[Dict[str, Any]] = None, finalize: bool = True,
) -> 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.
@ -122,14 +125,15 @@ def visualize_slice(self,
pyplot.show() pyplot.show()
def visualize_isosurface(self, def visualize_isosurface(
cell_data: numpy.ndarray, self,
level: Optional[float] = None, cell_data: NDArray,
which_shifts: int = 0, level: Optional[float] = None,
sample_period: int = 1, which_shifts: int = 0,
show_edges: bool = True, sample_period: int = 1,
finalize: bool = True, show_edges: bool = True,
) -> None: finalize: bool = True,
) -> None:
""" """
Draw an isosurface plot of the device. Draw an isosurface plot of the device.

4
gridlock/test/test_grid.py
Unescape Escape View File

@ -1,6 +1,6 @@
import pytest # type: ignore import pytest # type: ignore
import numpy # type: ignore import numpy
from numpy.testing import assert_allclose, assert_array_equal # type: ignore from numpy.testing import assert_allclose, assert_array_equal
from .. import Grid from .. import Grid

Write Preview
Loading…
Cancel
Save