[PolyCollection] rework PolyCollection into ndarrays of vertices and offsets

masque/shapes/poly_collection.py

@@ -1,97 +1,103 @@
-from typing import Any, cast, Iterable
+from typing import Any, cast, Self
-from collections.abc import Sequence
+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 rotation_matrix_2d, annotations_lt, annotations_eq, rep2key, annotations_t
-from ..utils import remove_colinear_vertices, remove_duplicate_vertices, 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
+    A collection of polygons, consisting of concatenated vertex arrays (N_m x 2 ndarray) which specify
-       implicitly-closed boundaries, and an offset.
+       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
+    A `normalized_form(...)` is available, but is untested and probably fairly slow.
-      passed vertex coordinates.
-
-    A `normalized_form(...)` is available, but can be quite slow with lots of vertices.
     """
     __slots__ = (
         '_vertex_lists',
+        '_vertex_offsets',
         # Inherited
         '_offset', '_repetition', '_annotations',
         )
 
-    _vertex_lists: list[NDArray[numpy.float64]]
+    _vertex_lists: NDArray[numpy.float64]
-    """ List of ndarrays (N_m x 2)  of vertices `[ [[x0, y0], [x1, y1], ...] ]` """
+    """ 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], ...] ]`
+        Vertices of the polygons, ((N+M+...) x 2). Use with `vertex_offsets`.
-
-        When setting, note that a copy will be made,
         """
         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 isinstance(vertex_lists, numpy.ndarray)
-            assert all(isinstance(vv, numpy.ndarray) for vv in vertex_lists)
+            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.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)
+        other = cast('PolyCollection', other)
-        for vv, oo in zip(self.vertices, other.vertices):
+
+        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,
+            vertices = vv,
-            offset=self.offset,
+            offset = self.offset,
-            repetition=self.repetition.copy(),
+            repetition = self.repetition.copy(),
-            annotations=copy.deepcopy(self.annotations),
+            annotations = copy.deepcopy(self.annotations),
-            ) for vv in self.vertex_lists]
+            ) 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]
+        return numpy.vstack((self.offset + numpy.min(self._vertex_lists, axis=0),
-        maxs = [numpy.max(vv, axis=0) for vv in self.vertex_lists]
+                             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:
+            rot = rotation_matrix_2d(theta)
-                vv[:] = numpy.dot(rotation_matrix_2d(theta), vv.T).T
+            self._vertex_lists = numpy.einsum('ij,kj->ki', rot, self._vertex_lists_)
         return self
 
-    def mirror(self, axis: int = 0) -> 'Polygon':
+    def mirror(self, axis: int = 0) -> Self:
-        for vv in self.vertex_lists:
+        self._vertex_lists[:, axis - 1] *= -1
-            vv[:, axis - 1] *= -1
         return self
 
-    def scale_by(self, c: float) -> 'Polygon':
+    def scale_by(self, c: float) -> Self:
-        for vv in self.vertex_lists:
+        self.vertex_lists *= c
-            vv *= 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)
+        meanv = self._vertex_lists.mean(axis=0)
-        zeroed_vertices = [vv - meanv for vv in self.vertex_lists]
+        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)
+        rotated_vertices = numpy.einsum('ij,kj->ki', rotation_matrix_2d(-rotation), normed_vertices)
-                                        for v in normed_vertices])
 
-        # Reorder the vertices so that the one with lowest x, then y, comes first.
+        # TODO consider how to reorder vertices for polycollection
-        x_min = rotated_vertices[:, 0].argmin()
+        ## Reorder the vertices so that the one with lowest x, then y, comes first.
-        if not is_scalar(x_min):
+        #x_min = rotated_vertices[:, 0].argmin()
-            y_min = rotated_vertices[x_min, 1].argmin()
+        #if not is_scalar(x_min):
-            x_min = cast(Sequence, x_min)[y_min]
+        #    y_min = rotated_vertices[x_min, 1].argmin()
-        reordered_vertices = numpy.roll(rotated_vertices, -x_min, axis=0)
+        #    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)
+        centroid = self.offset + self.vertex_lists.mean(axis=0)
-        return f'<PolyCollection centroid {centroid} p{len(self.vertex_lists)}>'
+        return f'<PolyCollection centroid {centroid} p{len(self.vertex_offsets)}>'