From 6db836d3a3c8a972079cfdad76ae65d0378da067 Mon Sep 17 00:00:00 2001 From: jan Date: Tue, 10 Oct 2017 16:42:11 -0700 Subject: [PATCH] Split Grid class into multiple files --- gridlock/draw.py | 349 +++++++++++++++++++++++++ gridlock/grid.py | 598 +------------------------------------------ gridlock/position.py | 105 ++++++++ gridlock/read.py | 181 +++++++++++++ 4 files changed, 641 insertions(+), 592 deletions(-) create mode 100644 gridlock/draw.py create mode 100644 gridlock/position.py create mode 100644 gridlock/read.py diff --git a/gridlock/draw.py b/gridlock/draw.py new file mode 100644 index 0000000..1debf0a --- /dev/null +++ b/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) + diff --git a/gridlock/grid.py b/gridlock/grid.py index dc30826..2097ba6 100644 --- a/gridlock/grid.py +++ b/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() diff --git a/gridlock/position.py b/gridlock/position.py new file mode 100644 index 0000000..ae58b8b --- /dev/null +++ b/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 diff --git a/gridlock/read.py b/gridlock/read.py new file mode 100644 index 0000000..ee14879 --- /dev/null +++ b/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()