[PolyCollection] rework PolyCollection into ndarrays of vertices and offsets

This commit is contained in:
Jan Petykiewicz 2025-04-24 23:19:01 -07:00
parent d0c1b00d7e
commit 5368fd4e16

View File

@ -1,97 +1,103 @@
from typing import Any, cast, Iterable
from collections.abc import Sequence
from typing import Any, cast, Self
from collections.abc import Iterator
import copy
import functools
from itertools import chain
import numpy
from numpy import pi
from numpy.typing import NDArray, ArrayLike
from . import Shape, normalized_shape_tuple
from ..error import PatternError
from .polygon import Polygon
from ..repetition import Repetition
from ..utils import is_scalar, rotation_matrix_2d, annotations_lt, annotations_eq, rep2key
from ..utils import remove_colinear_vertices, remove_duplicate_vertices, annotations_t
from ..utils import rotation_matrix_2d, annotations_lt, annotations_eq, rep2key, annotations_t
@functools.total_ordering
class PolyCollection(Shape):
"""
A collection of polygons, consisting of list of vertex arrays (N_m x 2 ndarrays) which specify
implicitly-closed boundaries, and an offset.
A collection of polygons, consisting of concatenated vertex arrays (N_m x 2 ndarray) which specify
implicitly-closed boundaries, and an array of offets specifying the first vertex of each
successive polygon.
Note that the setter for `PolyCollection.vertex_list` creates a copy of the
passed vertex coordinates.
A `normalized_form(...)` is available, but can be quite slow with lots of vertices.
A `normalized_form(...)` is available, but is untested and probably fairly slow.
"""
__slots__ = (
'_vertex_lists',
'_vertex_offsets',
# Inherited
'_offset', '_repetition', '_annotations',
)
_vertex_lists: list[NDArray[numpy.float64]]
""" List of ndarrays (N_m x 2) of vertices `[ [[x0, y0], [x1, y1], ...] ]` """
_vertex_lists: NDArray[numpy.float64]
""" 2D NDArray ((N+M+...) x 2) of vertices `[[xa0, ya0], [xa1, ya1], ..., [xb0, yb0], [xb1, yb1], ... ]` """
_vertex_offsets: NDArray[numpy.intp]
""" 1D NDArray specifying the starting offset for each polygon """
# vertex_lists property
@property
def vertex_lists(self) -> Any: # mypy#3004 NDArray[numpy.float64]:
"""
Vertices of the polygons (ist of ndarrays (N_m x 2) `[ [[x0, y0], [x1, y1], ...] ]`
When setting, note that a copy will be made,
Vertices of the polygons, ((N+M+...) x 2). Use with `vertex_offsets`.
"""
return self._vertex_lists
@vertex_lists.setter
def vertex_lists(self, val: ArrayLike) -> None:
val = [numpy.array(vv, dtype=float) for vv in val]
for ii, vv in enumerate(val):
if len(vv.shape) < 2 or vv.shape[1] != 2:
raise PatternError(f'vertex_lists contents must be an Nx2 arrays (polygon #{ii} fails)')
if vv.shape[0] < 3:
raise PatternError(f'vertex_lists contents must have at least 3 vertices (Nx2 where N>2) (polygon ${ii} has shape {vv.shape})')
self._vertices = val
# xs property
@property
def xs(self) -> NDArray[numpy.float64]:
def vertex_offsets(self) -> Any: # mypy#3004 NDArray[numpy.intp]:
"""
All vertex x coords as a 1D ndarray
Starting offset (in `vertex_lists`) for each polygon
"""
return self.vertices[:, 0]
return self._vertex_offsets
@property
def vertex_slices(self) -> Iterator[slice]:
"""
Iterator which provides slices which index vertex_lists
"""
for ii, ff in chain(self._vertex_offsets, (self._vertex_lists.shape[0],)):
yield slice(ii, ff)
@property
def polygon_vertices(self) -> Iterator[NDArray[numpy.float64]]:
for slc in self.vertex_slices:
yield self._vertex_lists[slc]
def __init__(
self,
vertex_lists: Iterable[ArrayLike],
vertex_lists: ArrayLike,
vertex_offsets: ArrayLike,
*,
offset: ArrayLike = (0.0, 0.0),
rotation: float = 0.0,
repetition: Repetition | None = None,
annotations: annotations_t | None = None,
annotations: annotations_t = None,
raw: bool = False,
) -> None:
if raw:
assert isinstance(vertex_lists, list)
assert all(isinstance(vv, numpy.ndarray) for vv in vertex_lists)
assert isinstance(vertex_lists, numpy.ndarray)
assert isinstance(vertex_offsets, numpy.ndarray)
assert isinstance(offset, numpy.ndarray)
self._vertex_lists = vertex_lists
self._vertex_offsets = vertex_offsets
self._offset = offset
self._repetition = repetition
self._annotations = annotations if annotations is not None else {}
self._annotations = annotations
else:
self.vertices = vertices
self._vertex_lists = numpy.asarray(vertex_lists, dtype=float)
self._vertex_offsets = numpy.asarray(vertex_offsets, dtype=numpy.intp)
self.offset = offset
self.repetition = repetition
self.annotations = annotations if annotations is not None else {}
self.rotate(rotation)
self.annotations = annotations
if rotation:
self.rotate(rotation)
def __deepcopy__(self, memo: dict | None = None) -> 'PolyCollection':
def __deepcopy__(self, memo: dict | None = None) -> Self:
memo = {} if memo is None else memo
new = copy.copy(self)
new._offset = self._offset.copy()
new._vertex_lists = [vv.copy() for vv in self._vertex_lists]
new._vertex_lists = self._vertex_lists.copy()
new._vertex_offsets = self._vertex_offsets.copy()
new._annotations = copy.deepcopy(self._annotations)
return new
@ -99,7 +105,8 @@ class PolyCollection(Shape):
return (
type(self) is type(other)
and numpy.array_equal(self.offset, other.offset)
and all(numpy.array_equal(ss, oo) for ss, oo in zip(self.vertices, other.vertices))
and numpy.array_equal(self._vertex_lists, other._vertex_lists)
and numpy.array_equal(self._vertex_offsets, other._vertex_offsets)
and self.repetition == other.repetition
and annotations_eq(self.annotations, other.annotations)
)
@ -110,8 +117,9 @@ class PolyCollection(Shape):
return repr(type(self)) < repr(type(other))
return id(type(self)) < id(type(other))
other = cast(PolyCollection, other)
for vv, oo in zip(self.vertices, other.vertices):
other = cast('PolyCollection', other)
for vv, oo in zip(self.polygon_vertices, other.polygon_vertices, strict=False):
if not numpy.array_equal(vv, oo):
min_len = min(vv.shape[0], oo.shape[0])
eq_mask = vv[:min_len] != oo[:min_len]
@ -128,60 +136,41 @@ class PolyCollection(Shape):
return rep2key(self.repetition) < rep2key(other.repetition)
return annotations_lt(self.annotations, other.annotations)
def pop_as_polygon(self, index: int) -> 'Polygon':
"""
Remove one polygon from the list, and return it as a `Polygon` object.
Args:
index: which polygon to pop
"""
verts = self.vertex_lists.pop(index)
return Polygon(
vertices=verts,
offset=self.offset,
repetition=self.repetition.copy(),
annotations=copy.deepcopy(self.annotations),
)
def to_polygons(
self,
num_vertices: int | None = None, # unused # noqa: ARG002
max_arclen: float | None = None, # unused # noqa: ARG002
) -> list['Polygon']:
return [Polygon(
vertices=vv,
offset=self.offset,
repetition=self.repetition.copy(),
annotations=copy.deepcopy(self.annotations),
) for vv in self.vertex_lists]
vertices = vv,
offset = self.offset,
repetition = self.repetition.copy(),
annotations = copy.deepcopy(self.annotations),
) for vv in self.polygon_vertices]
def get_bounds_single(self) -> NDArray[numpy.float64]: # TODO note shape get_bounds doesn't include repetition
mins = [numpy.min(vv, axis=0) for vv in self.vertex_lists]
maxs = [numpy.max(vv, axis=0) for vv in self.vertex_lists]
return numpy.vstack((self.offset + numpy.min(self.vertex_lists, axis=0),
self.offset + numpy.max(self.vertex_lists, axis=0)))
return numpy.vstack((self.offset + numpy.min(self._vertex_lists, axis=0),
self.offset + numpy.max(self._vertex_lists, axis=0)))
def rotate(self, theta: float) -> 'Polygon':
def rotate(self, theta: float) -> Self:
if theta != 0:
for vv in self.vertex_lists:
vv[:] = numpy.dot(rotation_matrix_2d(theta), vv.T).T
rot = rotation_matrix_2d(theta)
self._vertex_lists = numpy.einsum('ij,kj->ki', rot, self._vertex_lists_)
return self
def mirror(self, axis: int = 0) -> 'Polygon':
for vv in self.vertex_lists:
vv[:, axis - 1] *= -1
def mirror(self, axis: int = 0) -> Self:
self._vertex_lists[:, axis - 1] *= -1
return self
def scale_by(self, c: float) -> 'Polygon':
for vv in self.vertex_lists:
vv *= c
def scale_by(self, c: float) -> Self:
self.vertex_lists *= c
return self
def normalized_form(self, norm_value: float) -> normalized_shape_tuple:
# Note: this function is going to be pretty slow for many-vertexed polygons, relative to
# other shapes
meanv = numpy.concatenate(self.vertex_lists).mean(axis=0)
zeroed_vertices = [vv - meanv for vv in self.vertex_lists]
meanv = self._vertex_lists.mean(axis=0)
zeroed_vertices = self._vertex_lists - [meanv]
offset = meanv + self.offset
scale = zeroed_vertices.std()
@ -189,22 +178,26 @@ class PolyCollection(Shape):
_, _, vertex_axis = numpy.linalg.svd(zeroed_vertices)
rotation = numpy.arctan2(vertex_axis[0][1], vertex_axis[0][0]) % (2 * pi)
rotated_vertices = numpy.vstack([numpy.dot(rotation_matrix_2d(-rotation), v)
for v in normed_vertices])
rotated_vertices = numpy.einsum('ij,kj->ki', rotation_matrix_2d(-rotation), normed_vertices)
# Reorder the vertices so that the one with lowest x, then y, comes first.
x_min = rotated_vertices[:, 0].argmin()
if not is_scalar(x_min):
y_min = rotated_vertices[x_min, 1].argmin()
x_min = cast(Sequence, x_min)[y_min]
reordered_vertices = numpy.roll(rotated_vertices, -x_min, axis=0)
# TODO consider how to reorder vertices for polycollection
## Reorder the vertices so that the one with lowest x, then y, comes first.
#x_min = rotated_vertices[:, 0].argmin()
#if not is_scalar(x_min):
# y_min = rotated_vertices[x_min, 1].argmin()
# x_min = cast('Sequence', x_min)[y_min]
#reordered_vertices = numpy.roll(rotated_vertices, -x_min, axis=0)
# TODO: normalize mirroring?
return ((type(self), reordered_vertices.data.tobytes()),
return ((type(self), rotated_vertices.data.tobytes() + self._vertex_offsets.tobytes()),
(offset, scale / norm_value, rotation, False),
lambda: Polygon(reordered_vertices * norm_value))
lambda: PolyCollection(
vertex_lists=rotated_vertices * norm_value,
vertex_offsets=self._vertex_offsets,
),
)
def __repr__(self) -> str:
centroid = self.offset + numpy.concatenate(self.vertex_lists).mean(axis=0)
return f'<PolyCollection centroid {centroid} p{len(self.vertex_lists)}>'
centroid = self.offset + self.vertex_lists.mean(axis=0)
return f'<PolyCollection centroid {centroid} p{len(self.vertex_offsets)}>'