From f4862099509d88ffd1917e2ca06e30ca39900283 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Fri, 8 Jan 2021 22:58:36 -0800 Subject: [PATCH] general cleanup and typing fixes --- gridlock/draw.py | 35 ++++++++++++++++-------------- gridlock/grid.py | 51 +++++++++++++++++++++++++------------------- gridlock/position.py | 2 +- 3 files changed, 49 insertions(+), 39 deletions(-) diff --git a/gridlock/draw.py b/gridlock/draw.py index 3023388..af80047 100644 --- a/gridlock/draw.py +++ b/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) diff --git a/gridlock/grid.py b/gridlock/grid.py index a3f4788..7dc62c6 100644 --- a/gridlock/grid.py +++ b/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) diff --git a/gridlock/position.py b/gridlock/position.py index 8d75608..282824f 100644 --- a/gridlock/position.py +++ b/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