typing and formatting updates

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,16 +16,18 @@ 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,
+def draw_polygons(
-                  cell_data: numpy.ndarray,
+        self,
+        cell_data: NDArray,
         surface_normal: int,
-                  center: numpy.ndarray,
+        center: ArrayLike,
-                  polygons: Sequence[numpy.ndarray],
+        polygons: Sequence[NDArray],
         thickness: float,
-                  foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t],
+        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_min = numpy.array([0, 0])
-    bd_2d_max = [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,13 +191,14 @@ def draw_polygons(self,
         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,
+        cell_data: NDArray,
         surface_normal: int,
-                 center: numpy.ndarray,
+        center: ArrayLike,
-                 polygon: numpy.ndarray,
+        polygon: ArrayLike,
         thickness: float,
-                 foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t],
+        foreground: Union[Sequence[foreground_t], foreground_t],
         ) -> None:
     """
     Draw a polygon on an axis-aligned plane.
@@ -212,12 +216,13 @@ def draw_polygon(self,
     self.draw_polygons(cell_data, surface_normal, center, [polygon], thickness, foreground)
 
 
-def draw_slab(self,
+def draw_slab(
-              cell_data: numpy.ndarray,
+        self,
+        cell_data: NDArray,
         surface_normal: int,
-              center: numpy.ndarray,
+        center: ArrayLike,
         thickness: float,
-              foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t],
+        foreground: Union[Sequence[foreground_t], foreground_t],
         ) -> None:
     """
     Draw an axis-aligned infinite slab.
@@ -262,11 +267,12 @@ def draw_slab(self,
     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,
+        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,14 +293,15 @@ def draw_cuboid(self,
     self.draw_polygon(cell_data, 2, center, p, thickness, foreground)
 
 
-def draw_cylinder(self,
+def draw_cylinder(
-                  cell_data: numpy.ndarray,
+        self,
+        cell_data: NDArray,
         surface_normal: int,
-                  center: numpy.ndarray,
+        center: ArrayLike,
         radius: float,
         thickness: float,
         num_points: int,
-                  foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t],
+        foreground: Union[Sequence[foreground_t], foreground_t],
         ) -> None:
     """
     Draw an axis-aligned cylinder. Approximated by a num_points-gon
@@ -314,9 +322,10 @@ def draw_cylinder(self,
     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,
+        rectangle: ArrayLike,
         direction: int,
         polarity: int,
         distance: float,
@@ -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)

gridlock/examples/ex0.py

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

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],
+    Yee_Shifts_E: ClassVar[NDArray] = 0.5 * numpy.array([
+        [1, 0, 0],
         [0, 1, 0],
-                                                               [0, 0, 1]], dtype=float)
+        [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],
+    Yee_Shifts_H: ClassVar[NDArray] = 0.5 * numpy.array([
+        [0, 1, 1],
         [1, 0, 1],
-                                                               [1, 1, 0]], dtype=float)
+        [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,9 +259,10 @@ class Grid:
         else:
             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],
+            shifts: ArrayLike = Yee_Shifts_E,
             periodic: Union[bool, Sequence[bool]] = False,
             ) -> None:
         """

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,
+def ind2pos(
-            ind: numpy.ndarray,
+        self,
+        ind: NDArray,
         which_shifts: Optional[int] = None,
         round_ind: bool = True,
         check_bounds: bool = True
-            ) -> numpy.ndarray:
+        ) -> 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,
+def pos2ind(
-            r: numpy.ndarray,
+        self,
+        r: ArrayLike,
         which_shifts: Optional[int],
         round_ind: bool = True,
         check_bounds: bool = True
-            ) -> numpy.ndarray:
+        ) -> 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.

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,
+def get_slice(
-              cell_data: numpy.ndarray,
+        self,
+        cell_data: NDArray,
         surface_normal: int,
         center: float,
         which_shifts: int = 0,
         sample_period: int = 1
-              ) -> numpy.ndarray:
+        ) -> NDArray:
     """
     Retrieve a slice of a grid.
     Interpolates if given a position between two planes.
@@ -75,8 +77,9 @@ def get_slice(self,
     return sliced_grid
 
 
-def visualize_slice(self,
+def visualize_slice(
-                    cell_data: numpy.ndarray,
+        self,
+        cell_data: NDArray,
         surface_normal: int,
         center: float,
         which_shifts: int = 0,
@@ -122,8 +125,9 @@ def visualize_slice(self,
         pyplot.show()
 
 
-def visualize_isosurface(self,
+def visualize_isosurface(
-                         cell_data: numpy.ndarray,
+        self,
+        cell_data: NDArray,
         level: Optional[float] = None,
         which_shifts: int = 0,
         sample_period: int = 1,

gridlock/test/test_grid.py

@@ -1,6 +1,6 @@
 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