diff --git a/masque/shapes/poly_collection.py b/masque/shapes/poly_collection.py index c91c4a3..5d2ba49 100644 --- a/masque/shapes/poly_collection.py +++ b/masque/shapes/poly_collection.py @@ -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'' + centroid = self.offset + self.vertex_lists.mean(axis=0) + return f''