add all files

.gitignore

@@ -0,0 +1,3 @@
+*.pyc
+__pycache__
+*.idea

README.md

@@ -0,0 +1,22 @@
+# Gridlock README
+
+Gridlock is a Python module for drawing on coupled grids.
+
+Gridlock is used primarily for 'painting' shapes in 3D on multiple grids which represent the
+same spatial region, but are offset from each other. It does straightforward natural <-> grid unit
+conversion and can handle non-uniform rectangular grids (the entire grid is generated based on
+the coordinates of the boundary points along each axis).
+
+## Installation
+
+Requirements:
+* python 3 (written and tested with 3.5)
+* numpy
+* [float_raster](https://mpxd.net/gogs/jan/float_raster)
+* matplotlib (optional, used for visualization functions)
+* mpl_toolkits.mplot3d (optional, used for isosurface visualization)
+* skimage (optional, used for isosurface visualization)
+
+Install with pip, via git:
+
+>pip install --upgrade git+https://mpxd.net/gogs/jan/gridlock.git

gridlock/__init__.py

@@ -0,0 +1,23 @@
+"""
+3D coupled grid generator
+
+Grid generator, used primarily for 'painting' shapes in 3D on multiple grids which represent the
+ same spatial region, but are offset from each other. It does straightforward natural <-> grid unit
+ conversion and can handle non-uniform rectangular grids (the entire grid is generated based on
+ the coordinates of the boundary points along each axis).
+
+Its primary purpose is for drawing Yee grids for electromagnetic simulations.
+
+
+Dependencies:
+- numpy
+- matplotlib            [Grid.visualize_*]
+- mpl_toolkits.mplot3d  [Grid.visualize_isosurface()]
+- skimage               [Grid.visualize_isosurface()]
+"""
+
+from .error import GridError
+from .direction import Direction
+from .grid import Grid
+
+__author__ = 'Jan Petykiewicz'

gridlock/_helpers.py

@@ -0,0 +1,10 @@
+from typing import Any
+
+
+def is_scalar(var: Any) -> bool:
+    """
+    Alias for 'not hasattr(var, "__len__")'
+
+    :param var: Checks if var has a length.
+    """
+    return not hasattr(var, "__len__")

gridlock/direction.py

@@ -0,0 +1,10 @@
+from enum import Enum
+
+
+class Direction(Enum):
+    """
+    Enum for axis->integer mapping
+    """
+    x = 0
+    y = 1
+    z = 2

gridlock/error.py

@@ -0,0 +1,9 @@
+class GridError(Exception):
+    """
+    Simple Exception for Grid
+    """
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        return repr(self.value)

gridlock/grid.py

@@ -0,0 +1,804 @@
+from typing import List, Tuple, Callable
+
+import numpy
+from numpy import diff, floor, ceil, zeros, hstack, newaxis
+
+import pickle
+import warnings
+
+from float_raster import raster
+
+# .visualize_* uses matplotlib
+# .visualize_isosurface uses skimage
+# .visualize_isosurface uses mpl_toolkits.mplot3d
+
+from . import GridError, Direction
+from ._helpers import is_scalar
+
+__author__ = 'Jan Petykiewicz'
+
+eps_callable_type = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray]
+
+
+class Grid(object):
+    """
+    Simulation grid generator intended for electromagnetic simulations.
+    Can be used to generate non-uniform rectangular grids (the entire grid
+    is generated based on the coordinates of the boundary points). Also does
+    straightforward natural <-> grid unit conversion.
+
+    self.grids[i][a,b,c] contains the value of epsilon for the cell located at
+          (xyz[0][a]+dxyz[0][a]*shifts[i, 0],
+           xyz[1][b]+dxyz[1][b]*shifts[i, 1],
+           xyz[2][c]+dxyz[2][c]*shifts[i, 2]).
+     You can get raw edge coordinates (exyz),
+                   center coordinates (xyz),
+                           cell sizes (dxyz),
+      from the properties named as above, or get them for a given grid by using the
+      self.shifted_*xyz(which_shifts) functions.
+
+     It is tricky to determine the size of the right-most cell after shifting,
+      since its right boundary should shift by shifts[i][a] * dxyz[a][dxyz[a].size],
+      where the dxyz element refers to a cell that does not exist.
+     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.
+    """
+
+    # Cell edges. Monotonically increasing without duplicates
+    exyz = []               # type: List[numpy.ndarray]
+
+    # epsilon (or mu, or whatever) grids
+    grids = []              # type: List[numpy.ndarray]
+
+    # [[x0 y0 z0], [x1, y1, z1], ...] offsets for grid 0,1,...
+    shifts = None           # type: numpy.ndarray
+
+    # For each axis, determines how far the rightmost boundary gets shifted
+    periodic = [False] * 3  # type: List[bool]
+
+    # Intended for use as a static constant
+    Yee_Shifts = 0.5 * numpy.array([[1, 0, 0],
+                                    [0, 1, 0],
+                                    [0, 0, 1],
+                                    [0, 1, 1],
+                                    [1, 0, 1],
+                                    [1, 1, 0]], dtype=float)    # type: numpy.ndarray
+
+    @property
+    def dxyz(self) -> List[numpy.ndarray]:
+        """
+        Cell sizes for each axis, no shifts applied
+
+        :return: List of 3 ndarrays of cell sizes
+        """
+        return [diff(self.exyz[a]) for a in range(3)]
+
+    @property
+    def xyz(self) -> List[numpy.ndarray]:
+        """
+        Cell centers for each axis, no shifts applied
+
+        :return: List of 3 ndarrays of cell edges
+        """
+        return [self.exyz[a][:-1] + self.dxyz[a] / 2.0 for a in range(3)]
+
+    @property
+    def shape(self) -> numpy.ndarray:
+        """
+        The number of cells in x, y, and z
+
+        :return: ndarray [x_centers.size, y_centers.size, z_centers.size]
+        """
+        return numpy.array([coord.size - 1 for coord in self.exyz], dtype=int)
+
+    @property
+    def dxyz_with_ghost(self) -> List[numpy.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
+         above to learn why this is necessary.
+
+         If periodic, final edge shifts same amount as first
+         Otherwise, final edge shifts same amount as second-to-last
+
+        :return: list of [dxs, dys, dzs] with each element same length as elements of self.xyz
+        """
+        el = [0 if p else -1 for p in self.periodic]
+        return [hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)]
+
+    @property
+    def center(self) -> numpy.ndarray:
+        """
+        Center position of the entire grid, no shifts applied
+        :return: ndarray [x_center, y_center, z_center]
+        """
+        # center is just average of first and last xyz, which is just the average of the
+        #  first two and last two exyz
+        centers = [(self.exyz[a][:1] + self.exyz[a][-1:]) / 4.0 for a in range(3)]
+        return numpy.array(centers, dtype=float)
+
+    @property
+    def dxyz_limits(self) -> Tuple[numpy.ndarray, numpy.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
+         weighted average is performed when shifting).
+
+        :return: List of 2 ndarrays, d_min=[min(dx), min(dy), min(dz)] and d_max=[...]
+        """
+        d_min = numpy.array([min(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
+
+    def shifted_exyz(self, which_shifts: int or None) -> List[numpy.ndarray]:
+        """
+        Returns edges for which_shifts.
+
+        :param which_shifts: Which grid (which shifts) to use, or None for unshifted
+        :return: List of 3 ndarrays of cell edges
+        """
+        if which_shifts is None:
+            return self.exyz
+        dxyz = self.dxyz_with_ghost
+        shifts = self.shifts[which_shifts, :]
+        return [self.exyz[a] + dxyz[a] * shifts[a] for a in range(3)]
+
+    def shifted_dxyz(self, which_shifts: int or None) -> List[numpy.ndarray]:
+        """
+        Returns cell sizes for which_shifts.
+
+        :param which_shifts: Which grid (which shifts) to use, or None for unshifted
+        :return: List of 3 ndarrays of cell sizes
+        """
+        if which_shifts is None:
+            return self.dxyz
+        shifts = self.shifts[which_shifts, :]
+        dxyz = self.dxyz_with_ghost
+        return [(dxyz[a][:-1] * (1 - shifts[a]) + dxyz[a][1:] * shifts[a]) for a in range(3)]
+
+    def shifted_xyz(self, which_shifts: int or None) -> List[numpy.ndarray]:
+        """
+        Returns cell centers for which_shifts.
+
+        :param which_shifts: Which grid (which shifts) to use, or None for unshifted
+        :return: List of 3 ndarrays of cell centers
+        """
+        if which_shifts is None:
+            return self.xyz
+        exyz = self.shifted_exyz(which_shifts)
+        dxyz = self.shifted_dxyz(which_shifts)
+        return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)]
+
+    def ind2pos(self,
+                ind: numpy.ndarray or List,
+                which_shifts: int or None=None,
+                round_ind: bool=True,
+                check_bounds: bool=True
+                ) -> numpy.ndarray:
+        """
+        Returns the natural position corresponding to the specified indices.
+         The resulting position is clipped to the bounds of the grid
+        (to cell centers if round_ind=True, or cell outer edges if round_ind=False)
+
+        :param ind: Indices of the position. Can be fractional. (3-element ndarray or list)
+        :param which_shifts: which grid number (shifts) to use
+        :param round_ind: Whether to round ind to the nearest integer position before indexing
+                (default True)
+        :param check_bounds: Whether to raise an GridError if the provided ind is outside of
+                the grid, as defined above (centers if round_ind, else edges) (default True)
+        :return: 3-element ndarray specifying the natural position
+        :raises: GridError
+        """
+        if which_shifts is not None and which_shifts >= self.shifts.shape[0]:
+            raise GridError('Invalid shifts')
+        ind = numpy.array(ind, dtype=float)
+
+        if check_bounds:
+            if round_ind:
+                low_bound = 0.0
+                high_bound = -1
+            else:
+                low_bound = -0.5
+                high_bound = -0.5
+            if (ind < low_bound).any() or (ind > self.shape - high_bound).any():
+                raise GridError('Position outside of grid: {}'.format(ind))
+
+        if round_ind:
+            rind = numpy.clip(numpy.round(ind), 0, self.shape - 1)
+            sxyz = self.shifted_xyz(which_shifts)
+            position = [sxyz[a][rind[a]].astype(int) for a in range(3)]
+        else:
+            sexyz = self.shifted_exyz(which_shifts)
+            position = [numpy.interp(ind[a], numpy.arange(sexyz[a].size) - 0.5, sexyz[a])
+                        for a in range(3)]
+        return numpy.array(position, dtype=float)
+
+    def pos2ind(self,
+                r: numpy.ndarray or List,
+                which_shifts: int or None,
+                round_ind: bool=True,
+                check_bounds: bool=True
+                ) -> numpy.ndarray:
+        """
+        Returns the indices corresponding to the specified natural position.
+             The resulting position is clipped to within the outer centers of the grid.
+
+        :param r: Natural position that we will convert into indices (3-element ndarray or list)
+        :param which_shifts: which grid number (shifts) to use
+        :param round_ind: Whether to round the returned indices to the nearest integers.
+        :param check_bounds: Whether to throw an GridError if r is outside the grid edges
+        :return: 3-element ndarray specifying the indices
+        :raises: GridError
+        """
+        r = numpy.squeeze(r)
+        if r.size != 3:
+            raise GridError('r must be 3-element vector: {}'.format(r))
+
+        if (which_shifts is not None) and (which_shifts >= self.shifts.shape[0]):
+            raise GridError('')
+
+        sexyz = self.shifted_exyz(which_shifts)
+
+        if check_bounds:
+            for a in range(3):
+                if self.shape[a] > 1 and (r[a] < sexyz[a][0] or r[a] > sexyz[a][-1]):
+                    raise GridError('Position[{}] outside of grid!'.format(a))
+
+        grid_pos = zeros((3,))
+        for a in range(3):
+            xi = numpy.digitize(r[a], sexyz[a])  # Figure out which cell we're in
+            xi_clipped = numpy.clip(xi, 1, sexyz[a].size - 1) - 1  # Clip back into grid bounds
+
+            # No need to interpolate if round_ind is true or we were outside the grid
+            if round_ind or xi != xi_clipped:
+                grid_pos[a] = xi_clipped
+            else:
+                # Interpolate
+                x = self.shifted_xyz(which_shifts)[a][xi]
+                dx = self.shifted_dxyz(which_shifts)[a][xi]
+                f = (r[a] - x) / dx
+                # Clip to centers
+                grid_pos[a] = numpy.clip(xi + f, 0, self.shape[a] - 1)
+        return grid_pos
+
+    def __init__(self,
+                 pixel_edge_coordinates: List[List or numpy.ndarray],
+                 shifts: numpy.ndarray or List=Yee_Shifts,
+                 initial: float or numpy.ndarray or List[float] or List[numpy.ndarray]=(1.0,)*3,
+                 num_grids: int=None,
+                 periodic: bool or List[bool]=False):
+        """
+        Initialize a new Grid
+
+        :param pixel_edge_coordinates: 3-element list of (ndarrays or lists) specifying the
+         coordinates of the pixel edges in each dimensions
+         (ie, [[x0, x1, x2,...], [y0,...], [z0,...]] where the first pixel has x-edges x=x0 and
+          x=x1, the second has edges x=x1 and x=x2, etc.)
+        :param shifts: Nx3 array containing [x, y, z] offsets for each of N grids.
+         Yee shifts are used by default.
+        :param initial: Grids are initialized to this value. If scalar, all grids are initialized
+         with ndarrays full of the scalar. If a list of scalars, grid[i] is initialized to an
+         ndarray full of initial[i]. If a list of ndarrays of the same shape as the grids, grid[i]
+         is set to initial[i]. Default 1.
+        :param num_grids: How many grids to create. Must be <= shifts.shape[0].
+         Default is shifts.shape[0]
+        :param periodic: Specifies how the sizes of edge cells are calculated; see main class
+         documentation. List of 3 bool, or a single bool that gets broadcast. Default False.
+        :raises: GridError
+        """
+        self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)]
+        for i in range(3):
+            if len(self.exyz[i]) != len(pixel_edge_coordinates[i]):
+                warnings.warn('Dimension {} had duplicate edge coordinates'.format(i))
+
+        if is_scalar(periodic):
+            periodic = [periodic] * 3
+        self.periodic = periodic
+
+        self.shifts = numpy.array(shifts, dtype=float)
+        if self.shifts.shape[1] != 3:
+            GridError('Misshapen shifts; second axis size should be 3,'
+                      ' shape is {}'.format(self.shifts.shape))
+
+        num_shifts = self.shifts.shape[0]
+        if num_grids is None:
+            num_grids = num_shifts
+        elif num_grids > num_shifts:
+            raise GridError('Number of grids exceeds number of shifts (%u)' % num_shifts)
+
+        grids_shape = hstack((num_grids, self.shape))
+        if is_scalar(initial):
+            self.grids = numpy.full(grids_shape, initial)
+        else:
+            if len(initial) < num_grids:
+                raise GridError('Too few initial grids specified!')
+
+            self.grids = [None] * num_grids
+            for i in range(num_grids):
+                if is_scalar(initial[i]):
+                    if initial[i] is not None:
+                        self.grids[i] = numpy.full(self.shape, initial[i])
+                else:
+                    if not numpy.array_equal(initial[i].shape, self.shape):
+                        raise GridError('Initial grid sizes must match given coordinates')
+                    self.grids[i] = initial[i]
+
+    @staticmethod
+    def load(filename: str) -> 'Grid':
+        """
+        Load a grid from a file
+
+        :param filename: Filename to load from.
+        """
+        with open(filename, 'rb') as f:
+            tmp_dict = pickle.load(f)
+
+        g = Grid([[-1, 1]] * 3)
+        g.__dict__.update(tmp_dict)
+        return g
+
+    def save(self, filename: str):
+        """
+        Save to file.
+
+        :param filename: Filename to save to.
+        """
+        with open(filename, 'wb') as f:
+            pickle.dump(self.__dict__, f, protocol=2)
+
+    def draw_polygons(self,
+                      surface_normal: Direction or int,
+                      center: List or numpy.ndarray,
+                      polygons: List[numpy.ndarray or List],
+                      thickness: float,
+                      eps: List[float or eps_callable_type] or float or eps_callable_type):
+        """
+        Draw polygons on an axis-aligned plane.
+
+        :param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or
+         integer in range(3)
+        :param center: 3-element ndarray or list specifying an offset applied to all the polygons
+        :param polygons: List of Nx2 or Nx3 ndarrays, each specifying the vertices of a polygon
+             (non-closed, clockwise). If Nx3, the surface_normal coordinate is ignored. Each polygon
+             must have at least 3 vertices.
+        :param thickness: Thickness of the layer to draw
+        :param eps: Value to draw with ('epsilon'). Can be scalar, callable, or a list
+             of any of these (1 per grid). Callable values should take ndarrays x, y, z of equal
+             shape and return an ndarray of equal shape containing the eps value at the given x, y,
+             and z (natural, not grid coordinates).
+        :raises: GridError
+        """
+        # Turn surface_normal into its integer representation
+        if isinstance(surface_normal, Direction):
+            surface_normal = surface_normal.value
+
+        if surface_normal not in range(3):
+            raise GridError('Invalid surface_normal direction')
+
+        center = numpy.squeeze(center)
+
+        # Check polygons, and remove redundant coordinates
+        surface = numpy.delete(range(3), surface_normal)
+
+        for i, polygon in enumerate(polygons):
+            malformed = 'Malformed polygon: (%i)' % i
+            if polygon.shape[1] not in (2, 3):
+                    raise GridError(malformed + 'must be a Nx2 or Nx3 ndarray')
+            if polygon.shape[1] == 3:
+                polygon = polygon[surface, :]
+
+            if not polygon.shape[0] > 2:
+                raise GridError(malformed + 'must consist of more than 2 points')
+            if polygon.ndim > 2 and not numpy.unique(polygon[:, surface_normal]).size == 1:
+                raise GridError(malformed + 'must be in plane with surface normal %s'
+                                % 'xyz'[surface_normal])
+
+        # Broadcast eps where necessary
+        if is_scalar(eps):
+            eps = [eps] * len(self.grids)
+        elif isinstance(eps, numpy.ndarray):
+            raise GridError('ndarray not supported for eps')
+
+        # ## 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]
+        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))
+        bd_min = numpy.insert(bd_2d_min, surface_normal, -thickness / 2.0) + center
+        bd_max = numpy.insert(bd_2d_max, surface_normal, +thickness / 2.0) + center
+
+        # 2) Find indices (bdi) just outside bd elements
+        buf = 2  # size of safety buffer
+        # Use s_min and s_max with unshifted pos2ind to get absolute limits on
+        #  the indices the polygons might affect
+        s_min = self.shifts.min(axis=0)
+        s_max = self.shifts.max(axis=0)
+        bdi_min = self.pos2ind(bd_min + s_min, None, round_ind=False, check_bounds=False) - buf
+        bdi_max = self.pos2ind(bd_max + s_max, None, round_ind=False, check_bounds=False) + buf
+        bdi_min = numpy.maximum(floor(bdi_min), 0).astype(int)
+        bdi_max = numpy.minimum(ceil(bdi_max), self.shape - 1).astype(int)
+
+        # 3) Adjust polygons for center
+        polygons = [poly + center[surface] for poly in polygons]
+
+        # iterate over grids
+        for (i, grid) in enumerate(self.grids):
+            # ## Evaluate or expand eps[i]
+            if callable(eps[i]):
+                # meshgrid over the (shifted) domain
+                domain = [self.shifted_xyz(i)[k][bdi_min[k]:bdi_max[k]+1] for k in range(3)]
+                (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():
+                    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):
+                return numpy.insert(vector, surface_normal, (val,))
+
+            w_xy = zeros((bdi_max - bdi_min + 1)[surface].astype(int))
+
+            # Draw each polygon separately
+            for polygon in polygons:
+
+                # Get the boundaries of the polygon
+                pbd_min = polygon.min(axis=0)
+                pbd_max = polygon.max(axis=0)
+
+                # Find indices in w_xy just outside polygon
+                #  using per-grid xy-weights (self.shifted_xyz())
+                corner_min = self.pos2ind(to_3d(pbd_min), i,
+                                          check_bounds=False)[surface].astype(int)
+                corner_max = self.pos2ind(to_3d(pbd_max), i,
+                                          check_bounds=False)[surface].astype(int)
+
+                # Find indices in w_xy which are modified by polygon
+                # First for the edge coordinates (+1 since we're indexing edges)
+                edge_slices = [numpy.s_[i:f + 2] for i, f in zip(corner_min, corner_max)]
+                # Then for the pixel centers (-bdi_min since we're
+                #  calculating weights within a subspace)
+                centers_slice = [numpy.s_[i:f + 1] for i, f in zip(corner_min - bdi_min[surface],
+                                                                   corner_max - bdi_min[surface])]
+
+                aa_x, aa_y = (self.shifted_exyz(i)[a][s] for a, s in zip(surface, edge_slices))
+                w_xy[centers_slice] += raster(polygon.T, aa_x, aa_y)
+
+            # Clamp overlapping polygons to 1
+            w_xy = numpy.minimum(w_xy, 1.0)
+
+            # 2) Generate weights in z-direction
+            w_z = numpy.zeros(((bdi_max - bdi_min + 1)[surface_normal], ))
+
+            def get_zi(offset):
+                pos_3d = to_3d([0, 0], center[surface_normal] + offset)
+                grid_coords = self.pos2ind(pos_3d, i, check_bounds=False, round_ind=False)
+                w_coord_fp = (grid_coords - bdi_min)[surface_normal]
+                w_coord = floor(w_coord_fp).astype(int)
+                return w_coord_fp, w_coord
+
+            zi_top_fp, zi_top = get_zi(+thickness/2.0)
+            zi_bot_fp, zi_bot = get_zi(-thickness/2.0)
+
+            w_z[zi_bot:zi_top + 1] = 1
+
+            if zi_top_fp != zi_top < self.shape[surface_normal] - 1:
+                f = zi_top_fp - zi_top
+                w_z[zi_top] = f
+            if zi_bot_fp != zi_bot > 0:
+                f = zi_bot_fp - zi_bot
+                w_z[zi_bot] = 1 - f
+
+            # 3) Generate total weight function
+            w = (w_xy[:, :, newaxis] * w_z).transpose(numpy.insert([0, 1], surface_normal, (2,)))
+
+            # ## Modify the grid
+            g_slice = [numpy.s_[bdi_min[a]:bdi_max[a] + 1] for a in range(3)]
+            self.grids[i][g_slice] = (1 - w) * self.grids[i][g_slice] + w * eps[i]
+
+    def draw_polygon(self,
+                     surface_normal: Direction or int,
+                     center: List or numpy.ndarray,
+                     polygon: List or numpy.ndarray,
+                     thickness: float,
+                     eps: List[float or eps_callable_type] or float or eps_callable_type):
+        """
+        Draw a polygon on an axis-aligned plane.
+
+        :param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or
+         integer in range(3)
+        :param center: 3-element ndarray or list specifying an offset applied to the polygon
+        :param polygon: Nx2 or Nx3 ndarray specifying the vertices of a polygon (non-closed,
+             clockwise). If Nx3, the surface_normal coordinate is ignored. Must have at least 3
+             vertices.
+        :param thickness: Thickness of the layer to draw
+        :param eps: Value to draw with ('epsilon'). See draw_polygons() for details.
+        """
+        self.draw_polygons(surface_normal, center, [polygon], thickness, eps)
+
+    def draw_slab(self,
+                  surface_normal: Direction or int,
+                  center: List or numpy.ndarray,
+                  thickness: float,
+                  eps: List[float or eps_callable_type] or float or eps_callable_type):
+        """
+        Draw an axis-aligned infinite slab.
+
+        :param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or
+         integer in range(3)
+        :param center: Surface_normal coordinate at the center of the slab
+        :param thickness: Thickness of the layer to draw
+        :param eps: Value to draw with ('epsilon'). See draw_polygons() for details.
+        """
+        # Turn surface_normal into its integer representation
+        if isinstance(surface_normal, Direction):
+            surface_normal = surface_normal.value
+        if surface_normal not in range(3):
+            raise GridError('Invalid surface_normal direction')
+
+        if not is_scalar(center):
+            center = numpy.squeeze(center)
+            if len(center) == 3:
+                center = center[surface_normal]
+            else:
+                raise GridError('Bad center: {}'.format(center))
+
+        # Find center of slab
+        center_shift = self.center
+        center_shift[surface_normal] = center
+
+        surface = numpy.delete(range(3), surface_normal)
+
+        xyz_min = numpy.array([self.xyz[a][0] for a in range(3)], dtype=float)[surface]
+        xyz_max = numpy.array([self.xyz[a][-1] for a in range(3)], dtype=float)[surface]
+
+        dxyz = numpy.array([max(self.dxyz[i]) for i in surface], dtype=float)
+
+        xyz_min -= 4 * dxyz
+        xyz_max += 4 * dxyz
+
+        p = numpy.array([[xyz_min[0], xyz_max[1]],
+                         [xyz_max[0], xyz_max[1]],
+                         [xyz_max[0], xyz_min[1]],
+                         [xyz_min[0], xyz_min[1]]], dtype=float)
+
+        self.draw_polygon(surface_normal, center_shift, p, thickness, eps)
+
+    # TODO: TEST ME!
+    def draw_cuboid(self,
+                    center: List or numpy.ndarray,
+                    dimensions: List or numpy.ndarray,
+                    eps: List[float or eps_callable_type] or float or eps_callable_type):
+        """
+        Draw an axis-aligned cuboid
+
+        :param center: 3-element ndarray or list specifying the cylinder's center
+        :param dimensions: 3-element list or ndarray containing the x, y, and z edge-to-edge
+            sizes of the cuboid
+        :param eps: Value to draw with ('epsilon'). See draw_polygons() for details.
+        """
+        p = numpy.array([[-dimensions[0], +dimensions[1]],
+                         [+dimensions[0], +dimensions[1]],
+                         [+dimensions[0], -dimensions[1]],
+                         [-dimensions[0], -dimensions[1]]], dtype=float)
+        thickness = dimensions[3]
+        self.draw_polygon(Direction.z, center, p, thickness, eps)
+
+    def draw_cylinder(self,
+                      surface_normal: Direction or int,
+                      center: List or numpy.ndarray,
+                      radius: float,
+                      thickness: float,
+                      num_points: int,
+                      eps: List[float or eps_callable_type] or float or eps_callable_type):
+        """
+        Draw an axis-aligned cylinder. Approximated by a num_points-gon
+
+        :param surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or
+         integer in range(3)
+        :param center: 3-element ndarray or list specifying the cylinder's center
+        :param radius: cylinder radius
+        :param thickness: Thickness of the layer to draw
+        :param num_points: The circle is approximated by a polygon with num_points vertices
+        :param eps: Value to draw with ('epsilon'). See draw_polygons() for details.
+        """
+        theta = numpy.linspace(0, 2*numpy.pi, num_points, endpoint=False)
+        x = radius * numpy.sin(theta)
+        y = radius * numpy.cos(theta)
+        polygon = hstack((x[:, newaxis], y[:, newaxis]))
+        self.draw_polygon(surface_normal, center, polygon, thickness, eps)
+
+    def draw_extrude_rectangle(self,
+                               rectangle: List or numpy.ndarray,
+                               direction: Direction or int,
+                               polarity: int,
+                               distance: float):
+        """
+        Extrude a rectangle of a previously-drawn structure along an axis.
+
+        :param rectangle: 2x3 ndarray or list specifying the rectangle's corners
+        :param direction: Direction to extrude in. Direction enum or int in range(3)
+        :param polarity: +1 or -1, direction along axis to extrude in
+        :param distance: How far to extrude
+        """
+        # 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')
+
+        s = numpy.sign(polarity)
+        surface = numpy.delete(range(3), direction)
+
+        rectangle = numpy.array(rectangle, dtype=float)
+        if s == 0:
+            raise GridError('0 is not a valid polarity')
+        if direction not in range(3):
+            raise GridError('Invalid direction: {}'.format(direction))
+        if rectangle[0, direction] != rectangle[1, direction]:
+            raise GridError('Rectangle entries along extrusion direction do not match.')
+
+        center = rectangle.sum(axis=0) / 2.0
+        center[direction] += s * distance / 2.0
+
+        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)
+        for i, grid in enumerate(self.grids):
+            z = self.pos2ind(rectangle[0, :], i, round_ind=False, check_bounds=False)[direction]
+
+            ind = [int(floor(z)) if i == direction else slice(None) for i in range(3)]
+
+            fpart = z - floor(z)
+            mult = [1-fpart, fpart][::s]  # reverses if s negative
+
+            eps = mult[0] * grid[ind]
+            ind[direction] += 1
+            eps += mult[1] * grid[ind]
+
+            def f_eps(xs, ys, zs):
+                # 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)
+                # reshape to original shape and keep only in-plane components
+                (qi, ri) = [numpy.reshape(xyzi[:, k], xs.shape) for k in surface]
+                return eps[qi, ri]
+
+            eps_func[i] = f_eps
+
+        self.draw_polygon(direction, center, p, thickness, eps_func)
+
+    def visualize_slice(self,
+                        surface_normal: Direction or int,
+                        center: float,
+                        which_shifts: int=0,
+                        sample_period: int=1):
+        """
+        Visualize a slice of a grid.
+        Interpolates if given a position between two planes.
+
+        :param surface_normal: Axis normal to the plane we're displaying. Can be a Direction or
+         integer in range(3)
+        :param center: Scalar specifying position along surface_normal axis.
+        :param which_shifts: Which grid to display. Default is the first grid (0).
+        :param sample_period: Period for down-sampling the image. Default 1 (disabled)
+        """
+        from matplotlib import pyplot
+
+        if not is_scalar(center) and numpy.isreal(center):
+            raise GridError('center must be a real scalar')
+
+        sp = round(sample_period)
+        if sp <= 0:
+            raise GridError('sample_period must be positive')
+
+        if not is_scalar(which_shifts) or which_shifts < 0:
+            raise GridError('Invalid which_shifts')
+
+        # Turn surface_normal into its integer representation
+        if isinstance(surface_normal, Direction):
+            surface_normal = surface_normal.value
+        if surface_normal not in range(3):
+            raise GridError('Invalid surface_normal direction')
+
+        surface = numpy.delete(range(3), surface_normal)
+
+        # Extract indices and weights of planes
+        center3 = numpy.insert([0, 0], surface_normal, (center,))
+        center_index = self.pos2ind(center3, which_shifts,
+                                    round_ind=False, check_bounds=False)[surface_normal]
+        centers = numpy.unique([floor(center_index), ceil(center_index)]).astype(int)
+        if len(centers) == 2:
+            fpart = center_index - floor(center_index)
+            w = [1-fpart, fpart]    # longer distance -> less weight
+        else:
+            w = [1]
+
+        c_min, c_max = (self.xyz[surface_normal][i] for i in [0, -1])
+        if center < c_min or center > c_max:
+            raise GridError('Coordinate of visualized plane must be within simulation domain')
+
+        # Extract grid values from planes above and below visualized slice
+        eps = zeros(self.shape[surface])
+        for ci, weight in zip(centers, w):
+            s = tuple(ci if a == surface_normal else numpy.s_[::sp] for a in range(3))
+            eps += weight * self.grids[which_shifts][tuple(s)]
+
+        # Remove extra dimensions
+        eps = numpy.squeeze(eps)
+
+        surface = numpy.delete(range(3), surface_normal)
+
+        x, y = (self.shifted_exyz(which_shifts)[a] for a in surface)
+        xmesh, ymesh = numpy.meshgrid(x, y, indexing='ij')
+        x_label, y_label = ('xyz'[a] for a in surface)
+
+        pyplot.figure()
+        pyplot.pcolormesh(xmesh, ymesh, eps)
+        pyplot.colorbar()
+        pyplot.gca().set_aspect('equal', adjustable='box')
+        pyplot.xlabel(x_label)
+        pyplot.ylabel(y_label)
+        pyplot.show()
+
+    def visualize_isosurface(self,
+                             level: float=None,
+                             which_shifts: int=0,
+                             sample_period: int=1,
+                             show_edges: bool=True):
+        """
+        Draw an isosurface plot of the device.
+
+        :param level: Value at which to find isosurface. Default (None) uses mean value in grid.
+        :param which_shifts: Which grid to display. Default is the first grid (0).
+        :param sample_period: Period for down-sampling the image. Default 1 (disabled)
+        :param show_edges: Whether to draw triangle edges. Default True
+        """
+        from matplotlib import pyplot
+        import skimage.measure
+        # Claims to be unused, but needed for subplot(projection='3d')
+        from mpl_toolkits.mplot3d import Axes3D
+
+        # Get data from self.grids
+        grid = self.grids[which_shifts][::sample_period, ::sample_period, ::sample_period]
+        if level is None:
+            level = grid.mean()
+
+        # Find isosurface with marching cubes
+        verts, faces = skimage.measure.marching_cubes(grid, level)
+
+        # Convert vertices from index to position
+        pos_verts = numpy.array([self.ind2pos(verts[i, :], which_shifts, round_ind=False)
+                                 for i in range(verts.shape[0])], dtype=float)
+        xs, ys, zs = (pos_verts[:, a] for a in range(3))
+
+        # Draw the plot
+        fig = pyplot.figure()
+        ax = fig.add_subplot(111, projection='3d')
+        if show_edges:
+            ax.plot_trisurf(xs, ys, faces, zs)
+        else:
+            ax.plot_trisurf(xs, ys, faces, zs, edgecolor='none')
+
+        # Add a fake plot of a cube to force the axes to be equal lengths
+        max_range = numpy.array([xs.max() - xs.min(),
+                                 ys.max() - ys.min(),
+                                 zs.max() - zs.min()], dtype=float).max()
+        mg = numpy.mgrid[-1:2:2, -1:2:2, -1:2:2]
+        xbs = 0.5 * max_range * mg[0].flatten() + 0.5 * (xs.max() + xs.min())
+        ybs = 0.5 * max_range * mg[1].flatten() + 0.5 * (ys.max() + ys.min())
+        zbs = 0.5 * max_range * mg[2].flatten() + 0.5 * (zs.max() + zs.min())
+        # Comment or uncomment following both lines to test the fake bounding box:
+        for xb, yb, zb in zip(xbs, ybs, zbs):
+            ax.plot([xb], [yb], [zb], 'w')
+
+        pyplot.show()

setup.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+
+from distutils.core import setup
+
+setup(name='masque',
+      version='0.1',
+      description='Lithography mask library',
+      author='Jan Petykiewicz',
+      author_email='anewusername@gmail.com',
+      url='https://mpxd.net/gogs/jan/masque',
+      packages=['masque'],
+     )
+