general cleanup and typing fixes

gridlock/draw.py

@@ -44,6 +44,7 @@ def draw_polygons(self,
     # Turn surface_normal into its integer representation
     if isinstance(surface_normal, Direction):
         surface_normal = surface_normal.value
+    assert(isinstance(surface_normal, int))
 
     if surface_normal not in range(3):
         raise GridError('Invalid surface_normal direction')
@@ -96,8 +97,13 @@ def draw_polygons(self,
     # 3) Adjust polygons for center
     polygons = [poly + center[surface] for poly in polygons]
 
+    # ## Generate weighing function
+    def to_3d(vector: numpy.ndarray, val: float = 0.0) -> numpy.ndarray:
+        v_2d = numpy.array(vector, dtype=float)
+        return numpy.insert(v_2d, surface_normal, (val,))
+
     # iterate over grids
-    for (i, grid) in enumerate(self.grids):
+    for i, grid in enumerate(self.grids):
         # ## Evaluate or expand eps[i]
         if callable(eps[i]):
             # meshgrid over the (shifted) domain
@@ -105,17 +111,14 @@ def draw_polygons(self,
             (x0, y0, z0) = numpy.meshgrid(*domain, indexing='ij')
 
             # evaluate on the meshgrid
-            eps[i] = eps[i](x0, y0, z0)
-            if not numpy.isfinite(eps[i]).all():
+            eps_i = eps[i](x0, y0, z0)
+            if not numpy.isfinite(eps_i).all():
                 raise GridError('Non-finite values in eps[%u]' % i)
         elif not is_scalar(eps[i]):
             raise GridError('Unsupported eps[{}]: {}'.format(i, type(eps[i])))
-        # do nothing if eps[i] is scalar non-callable
-
-        # ## Generate weighing function
-        def to_3d(vector: List or numpy.ndarray, val: float=0.0):
-            v_2d = numpy.array(vector, dtype=float)
-            return numpy.insert(v_2d, surface_normal, (val,))
+        else:
+            # eps[i] is scalar non-callable
+            eps_i = eps[i]
 
         w_xy = zeros((bdi_max - bdi_min + 1)[surface].astype(int))
 
@@ -183,7 +186,7 @@ def draw_polygons(self,
 
         # ## Modify the grid
         g_slice = (i,) + tuple(numpy.s_[bdi_min[a]:bdi_max[a] + 1] for a in range(3))
-        self.grids[g_slice] = (1 - w) * self.grids[g_slice] + w * eps[i]
+        self.grids[g_slice] = (1 - w) * self.grids[g_slice] + w * eps_i
 
 
 def draw_polygon(self,
@@ -327,11 +330,9 @@ def draw_extrude_rectangle(self,
     # Turn extrude_direction into its integer representation
     if isinstance(direction, Direction):
         direction = direction.value
-    if abs(direction) not in range(3):
-        raise GridError('Invalid extrude_direction')
+    assert(isinstance(direction, int))
 
     s = numpy.sign(polarity)
-    surface = numpy.delete(range(3), direction)
 
     rectangle = numpy.array(rectangle, dtype=float)
     if s == 0:
@@ -344,12 +345,14 @@ def draw_extrude_rectangle(self,
     center = rectangle.sum(axis=0) / 2.0
     center[direction] += s * distance / 2.0
 
+    surface = numpy.delete(range(3), direction)
+
     dim = numpy.fabs(diff(rectangle, axis=0).T)[surface]
     p = numpy.vstack((numpy.array([-1, -1, 1, 1], dtype=float) * dim[0]/2.0,
                       numpy.array([-1, 1, 1, -1], dtype=float) * dim[1]/2.0)).T
     thickness = distance
 
-    eps_func = [None] * len(self.grids)
+    eps_func = []
     for i, grid in enumerate(self.grids):
         z = self.pos2ind(rectangle[0, :], i, round_ind=False, check_bounds=False)[direction]
 
@@ -367,10 +370,10 @@ def draw_extrude_rectangle(self,
             xyzi = numpy.array([self.pos2ind(qrs, which_shifts=i)
                                 for qrs in zip(xs.flat, ys.flat, zs.flat)], dtype=int)
             # reshape to original shape and keep only in-plane components
-            (qi, ri) = [numpy.reshape(xyzi[:, k], xs.shape) for k in surface]
+            qi, ri = (numpy.reshape(xyzi[:, k], xs.shape) for k in surface)
             return eps[qi, ri]
 
-        eps_func[i] = f_eps
+        eps_func.append(f_eps)
 
     self.draw_polygon(direction, center, p, thickness, eps_func)
 

gridlock/grid.py

@@ -16,7 +16,7 @@ __author__ = 'Jan Petykiewicz'
 eps_callable_type = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray]
 
 
-class Grid(object):
+class Grid:
     """
     Simulation grid generator intended for electromagnetic simulations.
     Can be used to generate non-uniform rectangular grids (the entire grid
@@ -44,19 +44,32 @@ class Grid(object):
      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]
+    """Cell edges. Monotonically increasing without duplicates."""
 
-    Yee_Shifts_E = 0.5 * numpy.array([[1, 0, 0],
-                                      [0, 1, 0],
-                                      [0, 0, 1]], dtype=float)    # type: numpy.ndarray
+    grids: numpy.ndarray
+    """epsilon (or mu, or whatever) grids. shape is (num_grids, X, Y, Z)"""
+
+    periodic: List[bool]
+    """For each axis, determines how far the rightmost boundary gets shifted. """
+
+    shifts: numpy.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)
     """Default shifts for Yee grid E-field"""
 
-    Yee_Shifts_H = 0.5 * numpy.array([[0, 1, 1],
-                                      [1, 0, 1],
-                                      [1, 1, 0]], dtype=float)    # type: numpy.ndarray
+    Yee_Shifts_H: ClassVar[numpy.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 draw_polygons, draw_polygon, draw_slab, draw_cuboid, \
-           draw_cylinder, draw_extrude_rectangle
+    from .draw import (
+        draw_polygons, draw_polygon, draw_slab, draw_cuboid,
+        draw_cylinder, draw_extrude_rectangle,
+        )
     from .read import get_slice, visualize_slice, visualize_isosurface
     from .position import ind2pos, pos2ind
 
@@ -237,23 +250,17 @@ class Grid(object):
         Raises:
             `GridError` on invalid input
         """
-        self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)]     # type: List[numpy.ndarray]
-        """Cell edges. Monotonically increasing without duplicates."""
-
-        self.grids = None           # type: numpy.ndarray
-        """epsilon (or mu, or whatever) grids. shape is (num_grids, X, Y, Z)"""
+        self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)]
+        self.shifts = numpy.array(shifts, dtype=float)
 
         for i in range(3):
             if len(self.exyz[i]) != len(pixel_edge_coordinates[i]):
                 warnings.warn('Dimension {} had duplicate edge coordinates'.format(i), stacklevel=2)
 
-        if is_scalar(periodic):
-            periodic = [periodic] * 3
-        self.periodic = periodic       # type: List[bool]
-        """For each axis, determines how far the rightmost boundary gets shifted. """
-
-        self.shifts = numpy.array(shifts, dtype=float)      # type: numpy.ndarray
-        """Offsets `[[x0, y0, z0], [x1, y1, z1], ...]` for grid `0,1,...`"""
+        if isinstance(periodic, bool):
+            self.periodic = [periodic] * 3
+        else:
+            self.periodic = list(periodic)
 
         if len(self.shifts.shape) != 2:
             raise GridError('Misshapen shifts: shifts must have two axes! '
@@ -276,7 +283,7 @@ class Grid(object):
             raise GridError('Number of grids exceeds number of shifts (%u)' % num_shifts)
 
         grids_shape = hstack((num_grids, self.shape))
-        if is_scalar(initial):
+        if isinstance(initial, (float, int)):
             if isinstance(initial, int):
                 warnings.warn('Initial value is an int, grids will be integer-typed!', stacklevel=2)
             self.grids = numpy.full(grids_shape, initial)

gridlock/position.py

@@ -42,7 +42,7 @@ def ind2pos(self,
     if check_bounds:
         if round_ind:
             low_bound = 0.0
-            high_bound = -1
+            high_bound = -1.0
         else:
             low_bound = -0.5
             high_bound = -0.5