""" Drawing-related methods for Grid class """ from typing import List, Optional, Union, Sequence, Callable import numpy # type: ignore from numpy import diff, floor, ceil, zeros, hstack, newaxis from float_raster import raster from . import GridError, Direction from ._helpers import is_scalar eps_callable_t = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray] def draw_polygons(self, surface_normal: Union[Direction, int], center: numpy.ndarray, polygons: Sequence[numpy.ndarray], thickness: float, eps: Union[Sequence[Union[float, eps_callable_t]], float, eps_callable_t], ) -> None: """ Draw polygons on an axis-aligned plane. Args: surface_normal: Axis normal to the plane we're drawing on. Can be a `Direction` or integer in `range(3)` center: 3-element ndarray or list specifying an offset applied to all the polygons 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. thickness: Thickness of the layer to draw eps: Value to draw with ('epsilon'). Can be scalar, callable, or a list of any of these (1 per grid). Callable values should take an ndarray the shape of the grid 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): v_2d = numpy.array(vector, dtype=float) return numpy.insert(v_2d, 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 = tuple(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): edges = self.shifted_exyz(i)[surface_normal] point = center[surface_normal] + offset grid_coord = numpy.digitize(point, edges) - 1 w_coord = grid_coord - bdi_min[surface_normal] if w_coord < 0: w_coord = 0 f = 0 elif w_coord >= w_z.size: w_coord = w_z.size - 1 f = 1 else: dz = self.shifted_dxyz(i)[surface_normal][grid_coord] f = (point - edges[grid_coord]) / dz return f, w_coord zi_top_f, zi_top = get_zi(+thickness / 2.0) zi_bot_f, zi_bot = get_zi(-thickness / 2.0) w_z[zi_bot + 1:zi_top] = 1 if zi_bot < zi_top: w_z[zi_top] = zi_top_f w_z[zi_bot] = 1 - zi_bot_f else: w_z[zi_bot] = zi_top_f - zi_bot_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 = (i,) + tuple(numpy.s_[bdi_min[a]:bdi_max[a] + 1] for a in range(3)) self.grids[g_slice] = (1 - w) * self.grids[g_slice] + w * eps[i] def draw_polygon(self, surface_normal: Union[Direction, int], center: numpy.ndarray, polygon: numpy.ndarray, thickness: float, eps: Union[Sequence[Union[float, eps_callable_t]], float, eps_callable_t], ) -> None: """ Draw a polygon on an axis-aligned plane. Args: surface_normal: Axis normal to the plane we're drawing on. Can be a Direction or integer in range(3) center: 3-element ndarray or list specifying an offset applied to the polygon 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. thickness: Thickness of the layer to draw 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: Union[Direction, int], center: numpy.ndarray, thickness: float, eps: Union[List[Union[float, eps_callable_t]], float, eps_callable_t], ) -> None: """ Draw an axis-aligned infinite slab. Args: surface_normal: Axis normal to the plane we're drawing on. Can be a `Direction` or integer in `range(3)` center: Surface_normal coordinate at the center of the slab thickness: Thickness of the layer to draw 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: numpy.ndarray, dimensions: numpy.ndarray, eps: Union[List[Union[float, eps_callable_t]], float, eps_callable_t], ) -> None: """ Draw an axis-aligned cuboid Args: center: 3-element ndarray or list specifying the cuboid's center dimensions: 3-element list or ndarray containing the x, y, and z edge-to-edge sizes of the cuboid 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: Union[Direction, int], center: numpy.ndarray, radius: float, thickness: float, num_points: int, eps: Union[List[Union[float, eps_callable_t]], float, eps_callable_t], ) -> None: """ Draw an axis-aligned cylinder. Approximated by a num_points-gon Args: surface_normal: Axis normal to the plane we're drawing on. Can be a `Direction` or integer in `range(3)` center: 3-element ndarray or list specifying the cylinder's center radius: cylinder radius thickness: Thickness of the layer to draw num_points: The circle is approximated by a polygon with `num_points` vertices 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: numpy.ndarray, direction: Union[Direction, int], polarity: int, distance: float, ) -> None: """ Extrude a rectangle of a previously-drawn structure along an axis. Args: rectangle: 2x3 ndarray or list specifying the rectangle's corners direction: Direction to extrude in. Direction enum or int in range(3) polarity: +1 or -1, direction along axis to extrude in 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)