gridlock/gridlock/draw.py

"""
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 = 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):
            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).clip(0)
            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 = (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: 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 cuboid'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)