Split Grid class into multiple files

gridlock/draw.py

@@ -0,0 +1,349 @@
+"""
+Drawing-related methods for Grid class
+"""
+from typing import List
+
+import numpy
+from numpy import diff, floor, ceil, zeros, hstack, newaxis
+
+from float_raster import raster
+
+from . import GridError, Direction
+from ._helpers import is_scalar
+
+
+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] + 0.5
+            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]:
+            f = zi_top_fp - zi_top
+            w_z[zi_top] = f
+        if zi_bot_fp != zi_bot > -1:
+            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)
+
+
+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) / 2.0
+    thickness = dimensions[2]
+    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)
+

gridlock/grid.py

@@ -7,15 +7,10 @@ import pickle
 import warnings
 import copy
 
-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]
@@ -68,6 +63,11 @@ class Grid(object):
                                       [1, 0, 1],
                                       [1, 1, 0]], dtype=float)    # type: numpy.ndarray
 
+    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
+
     @property
     def dxyz(self) -> List[numpy.ndarray]:
         """
@@ -200,98 +200,6 @@ class Grid(object):
             raise GridError('autoshifting requires exactly 3 grids')
         return [self.shifted_dxyz(which_shifts=a)[a] for a in range(3)]
 
-    def ind2pos(self,
-                ind: numpy.ndarray or List,
-                which_shifts: int = None,
-                round_ind: bool = True,
-                check_bounds: bool = True
-                ) -> numpy.ndarray:
-        """
-        Returns the natural position corresponding to the specified cell center 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).astype(int), 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 cell-center 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('Invalid which_shifts: {}'.format(which_shifts))
-
-        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]) - 1 # Figure out which cell we're in
-            xi_clipped = numpy.clip(xi, 0, sexyz[a].size - 2)  # 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],
@@ -400,497 +308,3 @@ class Grid(object):
         """
         return copy.deepcopy(self)
 
-    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] + 0.5
-                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]:
-                f = zi_top_fp - zi_top
-                w_z[zi_top] = f
-            if zi_bot_fp != zi_bot > -1:
-                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) / 2.0
-        thickness = dimensions[2]
-        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 get_slice(self,
-                  surface_normal: Direction or int,
-                  center: float,
-                  which_shifts: int = 0,
-                  sample_period: int = 1
-                  ) -> numpy.ndarray:
-        """
-            Retrieve 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)
-            :return Array containing the portion of the grid.
-        """
-        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 selected plane must be within simulation domain')
-
-        # Extract grid values from planes above and below visualized slice
-        sliced_grid = 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))
-            sliced_grid += weight * self.grids[which_shifts][tuple(s)]
-
-        # Remove extra dimensions
-        sliced_grid = numpy.squeeze(sliced_grid)
-
-        return sliced_grid
-
-    def visualize_slice(self,
-                        surface_normal: Direction or int,
-                        center: float,
-                        which_shifts: int = 0,
-                        sample_period: int = 1,
-                        finalize: bool = True,
-                        pcolormesh_args: Dict = None):
-        """
-        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)
-        :param finalize: Whether to call pyplot.show() after constructing the plot. Default True
-        """
-        from matplotlib import pyplot
-
-        # Set surface normal to its integer value
-        if isinstance(surface_normal, Direction):
-            surface_normal = surface_normal.value
-
-        if pcolormesh_args is None:
-            pcolormesh_args = {}
-
-        grid_slice = self.get_slice(surface_normal=surface_normal,
-                                    center=center,
-                                    which_shifts=which_shifts,
-                                    sample_period=sample_period)
-
-        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, grid_slice, **pcolormesh_args)
-        pyplot.colorbar()
-        pyplot.gca().set_aspect('equal', adjustable='box')
-        pyplot.xlabel(x_label)
-        pyplot.ylabel(y_label)
-        if finalize:
-            pyplot.show()
-
-    def visualize_isosurface(self,
-                             level: float = None,
-                             which_shifts: int = 0,
-                             sample_period: int = 1,
-                             show_edges: bool = True,
-                             finalize: 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
-        :param finalize: Whether to call pyplot.show() after constructing the plot. 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')
-
-        if finalize:
-            pyplot.show()

gridlock/position.py

@@ -0,0 +1,105 @@
+"""
+Position-related methods for Grid class
+"""
+
+from typing import List
+
+import numpy
+from numpy import zeros
+
+from . import GridError
+
+
+def ind2pos(self,
+            ind: numpy.ndarray or List,
+            which_shifts: int = None,
+            round_ind: bool = True,
+            check_bounds: bool = True
+            ) -> numpy.ndarray:
+    """
+    Returns the natural position corresponding to the specified cell center 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).astype(int), 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 cell-center 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('Invalid which_shifts: {}'.format(which_shifts))
+
+    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]) - 1 # Figure out which cell we're in
+        xi_clipped = numpy.clip(xi, 0, sexyz[a].size - 2)  # 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

gridlock/read.py

@@ -0,0 +1,181 @@
+"""
+Readback and visualization methods for Grid class
+"""
+from typing import Dict
+
+import numpy
+from numpy import floor, ceil, zeros
+
+from . import GridError, Direction
+from ._helpers import is_scalar
+
+# .visualize_* uses matplotlib
+# .visualize_isosurface uses skimage
+# .visualize_isosurface uses mpl_toolkits.mplot3d
+
+
+def get_slice(self,
+              surface_normal: Direction or int,
+              center: float,
+              which_shifts: int = 0,
+              sample_period: int = 1
+              ) -> numpy.ndarray:
+    """
+        Retrieve 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)
+        :return Array containing the portion of the grid.
+    """
+    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 selected plane must be within simulation domain')
+
+    # Extract grid values from planes above and below visualized slice
+    sliced_grid = 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))
+        sliced_grid += weight * self.grids[which_shifts][tuple(s)]
+
+    # Remove extra dimensions
+    sliced_grid = numpy.squeeze(sliced_grid)
+
+    return sliced_grid
+
+
+def visualize_slice(self,
+                    surface_normal: Direction or int,
+                    center: float,
+                    which_shifts: int = 0,
+                    sample_period: int = 1,
+                    finalize: bool = True,
+                    pcolormesh_args: Dict = None):
+    """
+    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)
+    :param finalize: Whether to call pyplot.show() after constructing the plot. Default True
+    """
+    from matplotlib import pyplot
+
+    # Set surface normal to its integer value
+    if isinstance(surface_normal, Direction):
+        surface_normal = surface_normal.value
+
+    if pcolormesh_args is None:
+        pcolormesh_args = {}
+
+    grid_slice = self.get_slice(surface_normal=surface_normal,
+                                center=center,
+                                which_shifts=which_shifts,
+                                sample_period=sample_period)
+
+    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, grid_slice, **pcolormesh_args)
+    pyplot.colorbar()
+    pyplot.gca().set_aspect('equal', adjustable='box')
+    pyplot.xlabel(x_label)
+    pyplot.ylabel(y_label)
+    if finalize:
+        pyplot.show()
+
+
+def visualize_isosurface(self,
+                         level: float = None,
+                         which_shifts: int = 0,
+                         sample_period: int = 1,
+                         show_edges: bool = True,
+                         finalize: 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
+    :param finalize: Whether to call pyplot.show() after constructing the plot. 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')
+
+    if finalize:
+        pyplot.show()