From ec5c77e018ff13b052bbbea851c92d95b527ed21 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Tue, 18 Oct 2022 19:44:30 -0700 Subject: [PATCH] typing and formatting updates --- gridlock/draw.py | 113 ++++++++++++++++++++----------------- gridlock/examples/ex0.py | 2 +- gridlock/grid.py | 58 ++++++++++--------- gridlock/position.py | 31 +++++----- gridlock/read.py | 54 ++++++++++-------- gridlock/test/test_grid.py | 4 +- 6 files changed, 142 insertions(+), 120 deletions(-) diff --git a/gridlock/draw.py b/gridlock/draw.py index 6385213..5390871 100644 --- a/gridlock/draw.py +++ b/gridlock/draw.py @@ -3,7 +3,8 @@ Drawing-related methods for Grid class """ 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 . import GridError @@ -15,17 +16,19 @@ from . import GridError # 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, - cell_data: numpy.ndarray, - surface_normal: int, - center: numpy.ndarray, - polygons: Sequence[numpy.ndarray], - thickness: float, - foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t], - ) -> None: +def draw_polygons( + self, + cell_data: NDArray, + surface_normal: int, + center: ArrayLike, + polygons: Sequence[NDArray], + thickness: float, + foreground: Union[Sequence[foreground_t], foreground_t], + ) -> None: """ Draw polygons on an axis-aligned plane. @@ -74,8 +77,8 @@ def draw_polygons(self, # ## Compute sub-domain of the grid occupied by polygons # 1) Compute outer bounds (bd) of polygons - bd_2d_min = [0, 0] - bd_2d_max = [0, 0] + bd_2d_min = numpy.array([0, 0]) + bd_2d_max = numpy.array([0, 0]) for polygon in polygons: bd_2d_min = numpy.minimum(bd_2d_min, polygon.min(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] # ## 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) 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 -def draw_polygon(self, - cell_data: numpy.ndarray, - surface_normal: int, - center: numpy.ndarray, - polygon: numpy.ndarray, - thickness: float, - foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t], - ) -> None: +def draw_polygon( + self, + cell_data: NDArray, + surface_normal: int, + center: ArrayLike, + polygon: ArrayLike, + thickness: float, + foreground: Union[Sequence[foreground_t], foreground_t], + ) -> None: """ 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) -def draw_slab(self, - cell_data: numpy.ndarray, - surface_normal: int, - center: numpy.ndarray, - thickness: float, - foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t], - ) -> None: +def draw_slab( + self, + cell_data: NDArray, + surface_normal: int, + center: ArrayLike, + thickness: float, + foreground: Union[Sequence[foreground_t], foreground_t], + ) -> None: """ 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) -def draw_cuboid(self, - cell_data: numpy.ndarray, - center: numpy.ndarray, - dimensions: numpy.ndarray, - foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t], - ) -> None: +def draw_cuboid( + self, + cell_data: NDArray, + center: ArrayLike, + dimensions: ArrayLike, + foreground: Union[Sequence[foreground_t], foreground_t], + ) -> None: """ Draw an axis-aligned cuboid @@ -278,6 +284,7 @@ def draw_cuboid(self, sizes of the cuboid 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]], [+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) -def draw_cylinder(self, - cell_data: numpy.ndarray, - surface_normal: int, - center: numpy.ndarray, - radius: float, - thickness: float, - num_points: int, - foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t], - ) -> None: +def draw_cylinder( + self, + cell_data: NDArray, + surface_normal: int, + center: ArrayLike, + radius: float, + thickness: float, + num_points: int, + foreground: Union[Sequence[foreground_t], foreground_t], + ) -> None: """ 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) -def draw_extrude_rectangle(self, - cell_data: numpy.ndarray, - rectangle: numpy.ndarray, - direction: int, - polarity: int, - distance: float, - ) -> None: +def draw_extrude_rectangle( + self, + cell_data: NDArray, + rectangle: ArrayLike, + direction: int, + polarity: int, + distance: float, + ) -> None: """ 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 foreground = mult[0] * grid[tuple(ind)] - ind[direction] += 1 + ind[direction] += 1 # type: ignore #(known safe) 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 xyzi = numpy.array([self.pos2ind(qrs, which_shifts=i) for qrs in zip(xs.flat, ys.flat, zs.flat)], dtype=int) diff --git a/gridlock/examples/ex0.py b/gridlock/examples/ex0.py index ca2ef55..59756bc 100644 --- a/gridlock/examples/ex0.py +++ b/gridlock/examples/ex0.py @@ -1,4 +1,4 @@ -import numpy # type: ignore +import numpy from gridlock import Grid diff --git a/gridlock/grid.py b/gridlock/grid.py index e320854..5f132a0 100644 --- a/gridlock/grid.py +++ b/gridlock/grid.py @@ -1,6 +1,7 @@ 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 import pickle @@ -10,7 +11,7 @@ import copy 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') @@ -48,23 +49,27 @@ class Grid: 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. """ - exyz: List[numpy.ndarray] + exyz: List[NDArray] """Cell edges. Monotonically increasing without duplicates.""" periodic: List[bool] """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,...`""" - Yee_Shifts_E: ClassVar[numpy.ndarray] = 0.5 * numpy.array([[1, 0, 0], - [0, 1, 0], - [0, 0, 1]], dtype=float) + Yee_Shifts_E: ClassVar[NDArray] = 0.5 * numpy.array([ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], dtype=float) """Default shifts for Yee grid E-field""" - Yee_Shifts_H: ClassVar[numpy.ndarray] = 0.5 * numpy.array([[0, 1, 1], - [1, 0, 1], - [1, 1, 0]], dtype=float) + Yee_Shifts_H: ClassVar[NDArray] = 0.5 * numpy.array([ + [0, 1, 1], + [1, 0, 1], + [1, 1, 0], + ], dtype=float) """Default shifts for Yee grid H-field""" from .draw import ( @@ -75,7 +80,7 @@ class Grid: from .position import ind2pos, pos2ind @property - def dxyz(self) -> List[numpy.ndarray]: + def dxyz(self) -> List[NDArray]: """ Cell sizes for each axis, no shifts applied @@ -85,7 +90,7 @@ class Grid: return [numpy.diff(ee) for ee in self.exyz] @property - def xyz(self) -> List[numpy.ndarray]: + def xyz(self) -> List[NDArray]: """ 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)] @property - def shape(self) -> numpy.ndarray: + def shape(self) -> NDArray[numpy.int_]: """ The number of cells in x, y, and z @@ -119,7 +124,7 @@ class Grid: return numpy.hstack((self.num_grids, self.shape)) @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 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)] @property - def center(self) -> numpy.ndarray: + def center(self) -> NDArray[numpy.float64]: """ Center position of the entire grid, no shifts applied @@ -148,7 +153,7 @@ class Grid: return numpy.array(centers, dtype=float) @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 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) 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. @@ -183,7 +188,7 @@ class Grid: 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`. @@ -210,7 +215,7 @@ class Grid: 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`. @@ -226,7 +231,7 @@ class Grid: dxyz = self.shifted_dxyz(which_shifts) 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. @@ -237,7 +242,7 @@ class Grid: raise GridError('Autoshifting requires exactly 3 grids') 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. @@ -254,11 +259,12 @@ class Grid: else: return numpy.full(self.cell_data_shape, fill_value, dtype=dtype) - def __init__(self, - pixel_edge_coordinates: Sequence[numpy.ndarray], - shifts: numpy.ndarray = Yee_Shifts_E, - periodic: Union[bool, Sequence[bool]] = False, - ) -> None: + def __init__( + self, + pixel_edge_coordinates: Sequence[ArrayLike], + shifts: ArrayLike = Yee_Shifts_E, + periodic: Union[bool, Sequence[bool]] = False, + ) -> None: """ Args: pixel_edge_coordinates: 3-element list of (ndarrays or lists) specifying the diff --git a/gridlock/position.py b/gridlock/position.py index 1224a12..35ec1df 100644 --- a/gridlock/position.py +++ b/gridlock/position.py @@ -1,19 +1,21 @@ """ 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 -def ind2pos(self, - ind: numpy.ndarray, - which_shifts: Optional[int] = None, - round_ind: bool = True, - check_bounds: bool = True - ) -> numpy.ndarray: +def ind2pos( + self, + ind: NDArray, + which_shifts: Optional[int] = None, + round_ind: bool = True, + check_bounds: bool = True + ) -> NDArray[numpy.float64]: """ Returns the natural position corresponding to the specified cell center indices. The resulting position is clipped to the bounds of the grid @@ -59,12 +61,13 @@ def ind2pos(self, return numpy.array(position, dtype=float) -def pos2ind(self, - r: numpy.ndarray, - which_shifts: Optional[int], - round_ind: bool = True, - check_bounds: bool = True - ) -> numpy.ndarray: +def pos2ind( + self, + r: ArrayLike, + which_shifts: Optional[int], + round_ind: bool = True, + check_bounds: bool = True + ) -> NDArray[numpy.float64]: """ Returns the cell-center indices corresponding to the specified natural position. The resulting position is clipped to within the outer centers of the grid. diff --git a/gridlock/read.py b/gridlock/read.py index aa059d5..055670c 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -3,7 +3,8 @@ Readback and visualization methods for Grid class """ from typing import Dict, Optional, Union, Any -import numpy # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike from . import GridError @@ -12,13 +13,14 @@ from . import GridError # .visualize_isosurface uses mpl_toolkits.mplot3d -def get_slice(self, - cell_data: numpy.ndarray, - surface_normal: int, - center: float, - which_shifts: int = 0, - sample_period: int = 1 - ) -> numpy.ndarray: +def get_slice( + self, + cell_data: NDArray, + surface_normal: int, + center: float, + which_shifts: int = 0, + sample_period: int = 1 + ) -> NDArray: """ Retrieve a slice of a grid. Interpolates if given a position between two planes. @@ -75,15 +77,16 @@ def get_slice(self, return sliced_grid -def visualize_slice(self, - cell_data: numpy.ndarray, - surface_normal: int, - center: float, - which_shifts: int = 0, - sample_period: int = 1, - finalize: bool = True, - pcolormesh_args: Optional[Dict[str, Any]] = None, - ) -> None: +def visualize_slice( + self, + cell_data: NDArray, + surface_normal: int, + center: float, + which_shifts: int = 0, + sample_period: int = 1, + finalize: bool = True, + pcolormesh_args: Optional[Dict[str, Any]] = None, + ) -> None: """ Visualize a slice of a grid. Interpolates if given a position between two planes. @@ -122,14 +125,15 @@ def visualize_slice(self, pyplot.show() -def visualize_isosurface(self, - cell_data: numpy.ndarray, - level: Optional[float] = None, - which_shifts: int = 0, - sample_period: int = 1, - show_edges: bool = True, - finalize: bool = True, - ) -> None: +def visualize_isosurface( + self, + cell_data: NDArray, + level: Optional[float] = None, + which_shifts: int = 0, + sample_period: int = 1, + show_edges: bool = True, + finalize: bool = True, + ) -> None: """ Draw an isosurface plot of the device. diff --git a/gridlock/test/test_grid.py b/gridlock/test/test_grid.py index fc54030..1e30bf3 100644 --- a/gridlock/test/test_grid.py +++ b/gridlock/test/test_grid.py @@ -1,6 +1,6 @@ import pytest # type: ignore -import numpy # type: ignore -from numpy.testing import assert_allclose, assert_array_equal # type: ignore +import numpy +from numpy.testing import assert_allclose, assert_array_equal from .. import Grid