402 lines
17 KiB
Python
402 lines
17 KiB
Python
"""
|
|
Drawing-related methods for Grid class
|
|
"""
|
|
from collections.abc import Sequence, Callable
|
|
|
|
import numpy
|
|
from numpy.typing import NDArray, ArrayLike
|
|
from float_raster import raster
|
|
|
|
from .utils import GridError, Slab, SlabDict, SlabProtocol, Extent, ExtentDict, ExtentProtocol
|
|
from .position import GridPosMixin
|
|
|
|
|
|
# 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?
|
|
|
|
|
|
foreground_callable_t = Callable[[NDArray, NDArray, NDArray], NDArray]
|
|
foreground_t = float | foreground_callable_t
|
|
|
|
|
|
class GridDrawMixin(GridPosMixin):
|
|
def draw_polygons(
|
|
self,
|
|
cell_data: NDArray,
|
|
foreground: Sequence[foreground_t] | foreground_t,
|
|
slab: SlabProtocol | SlabDict,
|
|
polygons: Sequence[ArrayLike],
|
|
*,
|
|
offset2d: ArrayLike = (0, 0),
|
|
) -> None:
|
|
"""
|
|
Draw polygons on an axis-aligned slab.
|
|
|
|
Args:
|
|
cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
|
|
foreground: Value to draw with ('brush color'). 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 foreground value at the given x, y,
|
|
and z (natural, not grid coordinates).
|
|
slab: `Slab` or slab-like dict specifying the slab in which the polygons will be drawn.
|
|
polygons: List of Nx2 or Nx3 ndarrays, each specifying the vertices of a polygon
|
|
(non-closed, clockwise). If Nx3, the `slab.axis`-th coordinate is ignored. Each
|
|
polygon must have at least 3 vertices.
|
|
offset2d: 2D offset to apply to polygon coordinates -- this offset is added directly
|
|
to the given polygon vertex coordinates. Default (0, 0).
|
|
|
|
Raises:
|
|
GridError
|
|
"""
|
|
if isinstance(slab, dict):
|
|
slab = Slab(**slab)
|
|
|
|
poly_list = [numpy.asarray(poly) for poly in polygons]
|
|
|
|
# Check polygons, and remove redundant coordinates
|
|
surface = numpy.delete(range(3), slab.axis)
|
|
|
|
for ii in range(len(poly_list)):
|
|
polygon = poly_list[ii]
|
|
malformed = f'Malformed polygon: ({ii})'
|
|
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, :]
|
|
poly_list[ii] = polygon
|
|
|
|
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[:, slab.axis]).size == 1:
|
|
raise GridError(malformed + 'must be in plane with surface normal ' + 'xyz'[slab.axis])
|
|
|
|
# Broadcast foreground where necessary
|
|
foregrounds: Sequence[foreground_callable_t] | Sequence[float]
|
|
if numpy.size(foreground) == 1: # type: ignore
|
|
foregrounds = [foreground] * len(cell_data) # type: ignore
|
|
elif isinstance(foreground, numpy.ndarray):
|
|
raise GridError('ndarray not supported for foreground')
|
|
else:
|
|
foregrounds = foreground # type: ignore
|
|
|
|
# ## Compute sub-domain of the grid occupied by polygons
|
|
# 1) Compute outer bounds (bd) of polygons
|
|
bd_2d_min = numpy.array([0, 0])
|
|
bd_2d_max = numpy.array([0, 0])
|
|
for polygon in poly_list:
|
|
bd_2d_min = numpy.minimum(bd_2d_min, polygon.min(axis=0)) + offset2d
|
|
bd_2d_max = numpy.maximum(bd_2d_max, polygon.max(axis=0)) + offset2d
|
|
bd_min = numpy.insert(bd_2d_min, slab.axis, slab.min)
|
|
bd_max = numpy.insert(bd_2d_max, slab.axis, slab.max)
|
|
|
|
# 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 offset2d
|
|
poly_list = [poly + offset2d for poly in poly_list]
|
|
|
|
# ## Generate weighing function
|
|
def to_3d(vector: NDArray, val: float = 0.0) -> NDArray[numpy.float64]:
|
|
v_2d = numpy.array(vector, dtype=float)
|
|
return numpy.insert(v_2d, slab.axis, (val,))
|
|
|
|
# iterate over grids
|
|
foreground_val: NDArray | float
|
|
for i, _ in enumerate(cell_data):
|
|
# ## Evaluate or expand foregrounds[i]
|
|
foregrounds_i = foregrounds[i]
|
|
if callable(foregrounds_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
|
|
foreground_val = foregrounds_i(x0, y0, z0)
|
|
if not numpy.isfinite(foreground_val).all():
|
|
raise GridError(f'Non-finite values in foreground[{i}]')
|
|
elif numpy.size(foregrounds_i) != 1:
|
|
raise GridError(f'Unsupported foreground[{i}]: {type(foregrounds_i)}')
|
|
else:
|
|
# foreground[i] is scalar non-callable
|
|
foreground_val = foregrounds_i
|
|
|
|
w_xy = numpy.zeros((bdi_max - bdi_min + 1)[surface].astype(int))
|
|
|
|
# Draw each polygon separately
|
|
for polygon in poly_list:
|
|
|
|
# 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, strict=True)]
|
|
# 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], strict=True))
|
|
|
|
aa_x, aa_y = (self.shifted_exyz(i)[a][s] for a, s in zip(surface, edge_slices, strict=True))
|
|
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)[slab.axis], ))
|
|
|
|
def get_zi(point: float, i=i, w_z=w_z) -> tuple[float, int]: # noqa: ANN001
|
|
edges = self.shifted_exyz(i)[slab.axis]
|
|
grid_coord = numpy.digitize(point, edges) - 1
|
|
w_coord = grid_coord - bdi_min[slab.axis]
|
|
|
|
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)[slab.axis][grid_coord]
|
|
f = (point - edges[grid_coord]) / dz
|
|
return f, w_coord
|
|
|
|
zi_top_f, zi_top = get_zi(slab.max)
|
|
zi_bot_f, zi_bot = get_zi(slab.min)
|
|
|
|
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], slab.axis, (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 * foreground_val
|
|
|
|
|
|
def draw_polygon(
|
|
self,
|
|
cell_data: NDArray,
|
|
foreground: Sequence[foreground_t] | foreground_t,
|
|
slab: SlabProtocol | SlabDict,
|
|
polygon: ArrayLike,
|
|
*,
|
|
offset2d: ArrayLike = (0, 0),
|
|
) -> None:
|
|
"""
|
|
Draw a polygon on an axis-aligned plane.
|
|
|
|
Args:
|
|
cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
|
|
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
|
|
slab: `Slab` or slab-like dict specifying the slab in which the polygon will be drawn.
|
|
polygon: Nx2 or Nx3 ndarray specifying the vertices of a polygon (non-closed,
|
|
clockwise). If Nx3, the `slab.axis`-th coordinate is ignored. Must have at
|
|
least 3 vertices.
|
|
offset2d: 2D offset to apply to polygon coordinates -- this offset is added directly
|
|
to the given polygon vertex coordinates. Default (0, 0).
|
|
"""
|
|
self.draw_polygons(
|
|
cell_data = cell_data,
|
|
slab = slab,
|
|
polygons = [polygon],
|
|
foreground = foreground,
|
|
offset2d = offset2d,
|
|
)
|
|
|
|
|
|
def draw_slab(
|
|
self,
|
|
cell_data: NDArray,
|
|
foreground: Sequence[foreground_t] | foreground_t,
|
|
slab: SlabProtocol | SlabDict,
|
|
) -> None:
|
|
"""
|
|
Draw an axis-aligned infinite slab.
|
|
|
|
Args:
|
|
cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
|
|
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
|
|
slab: `Slab` or slab-like dict (geometrical slab specification)
|
|
"""
|
|
if isinstance(slab, dict):
|
|
slab = Slab(**slab)
|
|
|
|
# Find center of slab
|
|
center_shift = self.center
|
|
center_shift[slab.axis] = slab.center
|
|
|
|
surface = numpy.delete(range(3), slab.axis)
|
|
u_min, u_max = self.exyz[surface[0]][[0, -1]]
|
|
v_min, v_max = self.exyz[surface[1]][[0, -1]]
|
|
|
|
margin = 4 * numpy.max([self.dxyz[surface[0]].max(),
|
|
self.dxyz[surface[1]].max()])
|
|
|
|
p = numpy.array([[u_min - margin, v_max + margin],
|
|
[u_max + margin, v_max + margin],
|
|
[u_max + margin, v_min - margin],
|
|
[u_min - margin, v_min - margin]], dtype=float)
|
|
|
|
self.draw_polygon(
|
|
cell_data = cell_data,
|
|
slab = slab,
|
|
polygon = p,
|
|
foreground = foreground,
|
|
)
|
|
|
|
|
|
def draw_cuboid(
|
|
self,
|
|
cell_data: NDArray,
|
|
foreground: Sequence[foreground_t] | foreground_t,
|
|
*,
|
|
x: ExtentProtocol | ExtentDict,
|
|
y: ExtentProtocol | ExtentDict,
|
|
z: ExtentProtocol | ExtentDict,
|
|
) -> None:
|
|
"""
|
|
Draw an axis-aligned cuboid
|
|
|
|
Args:
|
|
cell_data: Cell data to modify (e.g. created by `Grid.allocate()`)
|
|
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
|
|
x: `Extent` or extent-like dict specifying the x-extent of the cuboid.
|
|
y: `Extent` or extent-like dict specifying the y-extent of the cuboid.
|
|
z: `Extent` or extent-like dict specifying the z-extent of the cuboid.
|
|
"""
|
|
if isinstance(x, dict):
|
|
x = Extent(**x)
|
|
if isinstance(y, dict):
|
|
y = Extent(**y)
|
|
if isinstance(z, dict):
|
|
z = Extent(**z)
|
|
|
|
center = numpy.asarray([x.center, y.center, z.center])
|
|
|
|
p = numpy.array([[x.min, y.max],
|
|
[x.max, y.max],
|
|
[x.max, y.min],
|
|
[x.min, y.min]], dtype=float)
|
|
slab = Slab(axis=2, center=z.center, span=z.span)
|
|
self.draw_polygon(cell_data=cell_data, slab=slab, polygon=p, foreground=foreground)
|
|
|
|
|
|
def draw_cylinder(
|
|
self,
|
|
cell_data: NDArray,
|
|
h: SlabProtocol | SlabDict,
|
|
radius: float,
|
|
num_points: int,
|
|
center2d: ArrayLike,
|
|
foreground: Sequence[foreground_t] | foreground_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()`)
|
|
h:
|
|
radius:
|
|
num_points: The circle is approximated by a polygon with `num_points` vertices
|
|
center2d:
|
|
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
|
|
"""
|
|
if isinstance(h, dict):
|
|
h = Slab(**h)
|
|
|
|
theta = numpy.linspace(0, 2 * numpy.pi, num_points, endpoint=False)[:, None]
|
|
xy0 = numpy.hstack((numpy.sin(theta), numpy.cos(theta)))
|
|
polygon = radius * xy0
|
|
self.draw_polygon(cell_data=cell_data, slab=h, polygon=polygon, foreground=foreground, offset2d=center2d)
|
|
|
|
|
|
def draw_extrude_rectangle(
|
|
self,
|
|
cell_data: NDArray,
|
|
rectangle: ArrayLike,
|
|
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
|
|
"""
|
|
sgn = numpy.sign(polarity)
|
|
|
|
rectangle = numpy.asarray(rectangle, dtype=float)
|
|
if sgn == 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] += sgn * distance / 2.0
|
|
|
|
surface = numpy.delete(range(3), direction)
|
|
|
|
dim = numpy.fabs(numpy.diff(rectangle, axis=0).T)[surface]
|
|
poly = numpy.vstack((numpy.array([-1, -1, 1, 1], dtype=float) * dim[0] * 0.5,
|
|
numpy.array([-1, 1, 1, -1], dtype=float) * dim[1] * 0.5)).T
|
|
thickness = distance
|
|
|
|
foreground_func = []
|
|
for ii, grid in enumerate(cell_data):
|
|
zz = self.pos2ind(rectangle[0, :], ii, round_ind=False, check_bounds=False)[direction]
|
|
|
|
ind = [int(numpy.floor(zz)) if dd == direction else slice(None) for dd in range(3)]
|
|
|
|
fpart = zz - numpy.floor(zz)
|
|
mult = [1 - fpart, fpart][::sgn] # reverses if s negative
|
|
|
|
foreground = mult[0] * grid[tuple(ind)]
|
|
ind[direction] += 1 # type: ignore #(known safe)
|
|
foreground += mult[1] * grid[tuple(ind)]
|
|
|
|
def f_foreground(xs, ys, zs, ii=ii, foreground=foreground) -> NDArray[numpy.int64]: # noqa: ANN001
|
|
# transform from natural position to index
|
|
xyzi = numpy.array([self.pos2ind(qrs, which_shifts=ii)
|
|
for qrs in zip(xs.flat, ys.flat, zs.flat, strict=True)], dtype=numpy.int64)
|
|
# reshape to original shape and keep only in-plane components
|
|
qi, ri = (numpy.reshape(xyzi[:, kk], xs.shape) for kk in surface)
|
|
return foreground[qi, ri]
|
|
|
|
foreground_func.append(f_foreground)
|
|
|
|
slab = Slab(axis=direction, center=center[direction], span=thickness)
|
|
self.draw_polygon(cell_data, slab=slab, polygon=poly, foreground=foreground_func, offset2d=center[surface])
|
|
|