gridlock/gridlock/draw.py

"""
Drawing-related methods for Grid class
"""
from typing import List, Optional, Union, Sequence, Callable

import numpy        # type: ignore
from float_raster import raster

from . import GridError


# NOTE: Maybe it would make sense to create a GridDrawer class
#       which would hold both the `Grid` itself and `cell_data`
#       and could be used to call multiple `draw_*` methods
#       without having to pass `cell_data` again each time?


eps_callable_t = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray]


def draw_polygons(self,
                  cell_data: numpy.ndarray,
                  surface_normal: 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:
        cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
        surface_normal: Axis normal to the plane we're drawing on. 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
    """
    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 = f'Malformed polygon: ({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 '
                            + 'xyz'[surface_normal])

    # Broadcast eps where necessary
    if numpy.size(eps) == 1:
        eps = [eps] * len(cell_data)
    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(numpy.floor(bdi_min), 0).astype(int)
    bdi_max = numpy.minimum(numpy.ceil(bdi_max), self.shape - 1).astype(int)

    # 3) Adjust polygons for center
    polygons = [poly + center[surface] for poly in polygons]

    # ## Generate weighing function
    def to_3d(vector: numpy.ndarray, val: float = 0.0) -> numpy.ndarray:
        v_2d = numpy.array(vector, dtype=float)
        return numpy.insert(v_2d, surface_normal, (val,))

    # iterate over grids
    for i, grid in enumerate(cell_data):
        # ## 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(f'Non-finite values in eps[{i}]')
        elif numpy.size(eps[i]) != 1:
            raise GridError(f'Unsupported eps[{i}]: {type(eps[i])}')
        else:
            # eps[i] is scalar non-callable
            eps_i = eps[i]

        w_xy = numpy.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, i=i, w_z=w_z):
            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[:, :, None] * 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))
        cell_data[g_slice] = (1 - w) * cell_data[g_slice] + w * eps_i


def draw_polygon(self,
                 cell_data: numpy.ndarray,
                 surface_normal: 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:
        cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
        surface_normal: Axis normal to the plane we're drawing on. 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(cell_data, surface_normal, center, [polygon], thickness, eps)


def draw_slab(self,
              cell_data: numpy.ndarray,
              surface_normal: 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:
        cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
        surface_normal: Axis normal to the plane we're drawing on. Integer in `range(3)`.
        center: `surface_normal` coordinate value 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 surface_normal not in range(3):
        raise GridError('Invalid surface_normal direction')

    if numpy.size(center) != 1:
        center = numpy.squeeze(center)
        if len(center) == 3:
            center = center[surface_normal]
        else:
            raise GridError(f'Bad center: {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(cell_data, surface_normal, center_shift, p, thickness, eps)


def draw_cuboid(self,
                cell_data: numpy.ndarray,
                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:
        cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
        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(cell_data, 2, center, p, thickness, eps)


def draw_cylinder(self,
                  cell_data: numpy.ndarray,
                  surface_normal: 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:
        cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
        surface_normal: Axis normal to the plane we're drawing on. 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 = numpy.hstack((x[:, None], y[:, None]))
    self.draw_polygon(cell_data, surface_normal, center, polygon, thickness, eps)


def draw_extrude_rectangle(self,
                           cell_data: numpy.ndarray,
                           rectangle: numpy.ndarray,
                           direction: int,
                           polarity: int,
                           distance: float,
                           ) -> None:
    """
    Extrude a rectangle of a previously-drawn structure along an axis.

    Args:
        cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
        rectangle: 2x3 ndarray or list specifying the rectangle's corners
        direction: Direction to extrude in. Integer in `range(3)`.
        polarity: +1 or -1, direction along axis to extrude in
        distance: How far to extrude
    """
    s = numpy.sign(polarity)

    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(f'Invalid direction: {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

    surface = numpy.delete(range(3), direction)

    dim = numpy.fabs(numpy.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 = []
    for i, grid in enumerate(cell_data):
        z = self.pos2ind(rectangle[0, :], i, round_ind=False, check_bounds=False)[direction]

        ind = [int(numpy.floor(z)) if i == direction else slice(None) for i in range(3)]

        fpart = z - numpy.floor(z)
        mult = [1-fpart, fpart][::s]  # reverses if s negative

        eps = mult[0] * grid[tuple(ind)]
        ind[direction] += 1
        eps += mult[1] * grid[tuple(ind)]

        def f_eps(xs, ys, zs, i=i, eps=eps) -> numpy.ndarray:
            # 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.append(f_eps)

    self.draw_polygon(cell_data, direction, center, p, thickness, eps_func)