From febaaeff0bcce6e8de4ba94b22df462854de0ee6 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Wed, 18 Sep 2024 20:46:48 -0700 Subject: [PATCH] add Library functions for finding instances and extracting hierarchy added child_graph, parent_graph, child_order, find_refs_local and find_refs_global --- masque/library.py | 167 ++++++++++++++++++++++++++++++++++---- masque/ref.py | 10 +++ masque/utils/__init__.py | 1 + masque/utils/transform.py | 56 ++++++++++++- 4 files changed, 217 insertions(+), 17 deletions(-) diff --git a/masque/library.py b/masque/library.py index 001ae72..08a15b9 100644 --- a/masque/library.py +++ b/masque/library.py @@ -22,12 +22,13 @@ import copy from pprint import pformat from collections import defaultdict from abc import ABCMeta, abstractmethod +from graphlib import TopologicalSorter import numpy from numpy.typing import ArrayLike, NDArray from .error import LibraryError, PatternError -from .utils import rotation_matrix_2d, layer_t +from .utils import layer_t, apply_transforms from .shapes import Shape, Polygon from .label import Label from .abstract import Abstract @@ -474,24 +475,21 @@ class ILibraryView(Mapping[str, 'Pattern'], metaclass=ABCMeta): raise LibraryError(f'.dfs() called on pattern with circular reference to "{target}"') for ref in pattern.refs[target]: + ref_transforms: list[bool] | NDArray[numpy.float64] if transform is not False: - sign = numpy.ones(2) - if transform[3]: - sign[1] = -1 - xy = numpy.dot(rotation_matrix_2d(transform[2]), ref.offset * sign) - ref_transform = transform + (xy[0], xy[1], ref.rotation, ref.mirrored) - ref_transform[3] %= 2 + ref_transforms = apply_transforms(transform, ref.as_transforms()) else: - ref_transform = False + ref_transforms = [False] - self.dfs( - pattern=self[target], - visit_before=visit_before, - visit_after=visit_after, - hierarchy=hierarchy + (target,), - transform=ref_transform, - memo=memo, - ) + for ref_transform in ref_transforms: + self.dfs( + pattern=self[target], + visit_before=visit_before, + visit_after=visit_after, + hierarchy=hierarchy + (target,), + transform=ref_transform, + memo=memo, + ) if visit_after is not None: pattern = visit_after(pattern, hierarchy=hierarchy, memo=memo, transform=transform) @@ -510,6 +508,143 @@ class ILibraryView(Mapping[str, 'Pattern'], metaclass=ABCMeta): return self + def child_graph(self) -> dict[str, set[str | None]]: + """ + Return a mapping from pattern name to a set of all child patterns + (patterns it references). + + Returns: + Mapping from pattern name to a set of all pattern names it references. + """ + graph = {name: set(pat.refs.keys()) for name, pat in self.items()} + return graph + + def parent_graph(self) -> dict[str, set[str]]: + """ + Return a mapping from pattern name to a set of all parent patterns + (patterns which reference it). + + Returns: + Mapping from pattern name to a set of all patterns which reference it. + """ + igraph: dict[str, set[str]] = {name: set() for name in self} + for name, pat in self.items(): + for child, reflist in pat.refs.items(): + if reflist and child is not None: + igraph[child].add(name) + return igraph + + def child_order(self) -> list[str]: + """ + Return a topologically sorted list of all contained pattern names. + Child (referenced) patterns will appear before their parents. + + Return: + Topologically sorted list of pattern names. + """ + return list(TopologicalSorter(self.child_graph()).static_order()) + + def find_refs_local( + self, + name: str, + parent_graph: dict[str, set[str]] | None = None, + ) -> dict[str, list[NDArray[numpy.float64]]]: + """ + Find the location and orientation of all refs pointing to `name`. + Refs with a `repetition` are resolved into multiple instances (locations). + + Args: + name: Name of the referenced pattern. + parent_graph: Mapping from pattern name to the set of patterns which + reference it. Default (`None`) calls `self.parent_graph()`. + The provided graph may be for a superset of `self` (i.e. it may + contain additional patterns which are not present in self; they + will be ignored). + + Returns: + Mapping of {parent_name: transform_list}, where transform_list + is an Nx4 ndarray with rows + `(x_offset, y_offset, rotation_ccw_rad, mirror_across_x)`. + """ + instances = defaultdict(list) + if parent_graph is None: + parent_graph = self.parent_graph() + for parent in parent_graph[name]: + if parent not in self: # parent_graph may be a for a superset of self + continue + for ref in self[parent].refs[name]: + instances[parent].append(ref.as_transforms()) + + return instances + + def find_refs_global( + self, + name: str, + order: list[str] | None = None, + parent_graph: dict[str, set[str]] | None = None, + ) -> dict[tuple[str, ...], NDArray[numpy.float64]]: + """ + Find the absolute (top-level) location and orientation of all refs (including + repetitions) pointing to `name`. + + Args: + name: Name of the referenced pattern. + order: List of pattern names in which children are guaranteed + to appear before their parents (i.e. topologically sorted). + Default (`None`) calls `self.child_order()`. + parent_graph: Passed to `find_refs_local`. + Mapping from pattern name to the set of patterns which + reference it. Default (`None`) calls `self.parent_graph()`. + The provided graph may be for a superset of `self` (i.e. it may + contain additional patterns which are not present in self; they + will be ignored). + + Returns: + Mapping of `{hierarchy: transform_list}`, where `hierarchy` is a tuple of the form + `(toplevel_pattern, lvl1_pattern, ..., name)` and `transform_list` is an Nx4 ndarray + with rows `(x_offset, y_offset, rotation_ccw_rad, mirror_across_x)`. + """ + if name not in self: + return {} + if order is None: + order = self.child_order() + if parent_graph is None: + parent_graph = self.parent_graph() + + self_keys = set(self.keys()) + + transforms: dict[str, list[tuple[ + tuple[str, ...], + NDArray[numpy.float64] + ]]] + transforms = defaultdict(list) + for parent, vals in self.find_refs_local(name, parent_graph=parent_graph).items(): + transforms[parent] = [((name,), numpy.concatenate(vals))] + + for next_name in order: + if next_name not in transforms: + continue + if not parent_graph[next_name] & self_keys: + continue + + outers = self.find_refs_local(next_name, parent_graph=parent_graph) + inners = transforms.pop(next_name) + for parent, outer in outers.items(): + for path, inner in inners: + combined = apply_transforms(numpy.concatenate(outer), inner) + transforms[parent].append(( + (next_name,) + path, + combined, + )) + result = {} + for parent, targets in transforms.items(): + for path, instances in targets: + full_path = (parent,) + path + assert full_path not in result + result[full_path] = instances + return result + + class ILibrary(ILibraryView, MutableMapping[str, 'Pattern'], metaclass=ABCMeta): """ diff --git a/masque/ref.py b/masque/ref.py index b84114b..09e00b1 100644 --- a/masque/ref.py +++ b/masque/ref.py @@ -183,6 +183,16 @@ class Ref( self.rotation += pi return self + def as_transforms(self) -> NDArray[numpy.float64]: + xys = self.offset[None, :] + if self.repetition is not None: + xys = xys + self.repetition.displacements + transforms = numpy.empty((xys.shape[0], 4)) + transforms[:, :2] = xys + transforms[:, 2] = self.rotation + transforms[:, 3] = self.mirrored + return transforms + def get_bounds_single( self, pattern: 'Pattern', diff --git a/masque/utils/__init__.py b/masque/utils/__init__.py index ffa9e85..11391a5 100644 --- a/masque/utils/__init__.py +++ b/masque/utils/__init__.py @@ -24,6 +24,7 @@ from .transform import ( rotation_matrix_2d as rotation_matrix_2d, normalize_mirror as normalize_mirror, rotate_offsets_around as rotate_offsets_around, + apply_transforms as apply_transforms, ) from .comparisons import ( annotation2key as annotation2key, diff --git a/masque/utils/transform.py b/masque/utils/transform.py index 3071d99..e2ea9aa 100644 --- a/masque/utils/transform.py +++ b/masque/utils/transform.py @@ -5,7 +5,7 @@ from collections.abc import Sequence from functools import lru_cache import numpy -from numpy.typing import NDArray +from numpy.typing import NDArray, ArrayLike from numpy import pi @@ -57,8 +57,62 @@ def rotate_offsets_around( ) -> NDArray[numpy.float64]: """ Rotates offsets around a pivot point. + + Args: + offsets: Nx2 array, rows are (x, y) offsets + pivot: (x, y) location to rotate around + angle: rotation angle in radians + + Returns: + Nx2 ndarray of (x, y) position after the rotation is applied. """ offsets -= pivot offsets[:] = (rotation_matrix_2d(angle) @ offsets.T).T offsets += pivot return offsets + + +def apply_transforms( + outer: ArrayLike, + inner: ArrayLike, + tensor: bool = False, + ) -> NDArray[numpy.float64]: + """ + Apply a set of transforms (`outer`) to a second set (`inner`). + This is used to find the "absolute" transform for nested `Ref`s. + + The two transforms should be of shape Ox4 and Ix4. + Rows should be of the form `(x_offset, y_offset, rotation_ccw_rad, mirror_across_x)`. + The output will be of the form (O*I)x4 (if `tensor=False`) or OxIx4 (`tensor=True`). + + Args: + outer: Transforms for the container refs. Shape Ox4. + inner: Transforms for the contained refs. Shape Ix4. + tensor: If `True`, an OxIx4 array is returned, with `result[oo, ii, :]` corresponding + to the `oo`th `outer` transform applied to the `ii`th inner transform. + If `False` (default), this is concatenated into `(O*I)x4` to allow simple + chaining into additional `apply_transforms()` calls. + + Returns: + OxIx4 or (O*I)x4 array. Final dimension is + `(total_x, total_y, total_rotation_ccw_rad, net_mirrored_x)`. + """ + outer = numpy.atleast_2d(outer).astype(float, copy=False) + inner = numpy.atleast_2d(inner).astype(float, copy=False) + + # If mirrored, flip y's + xy_mir = numpy.tile(inner[:, :2], (outer.shape[0], 1, 1)) # dims are outer, inner, xyrm + xy_mir[outer[:, 3].astype(bool), :, 1] *= -1 + + rot_mats = [rotation_matrix_2d(angle) for angle in outer[:, 2]] + xy = numpy.einsum('ort,oit->oir', rot_mats, xy_mir) + + tot = numpy.empty((outer.shape[0], inner.shape[0], 4)) + tot[:, :, :2] = outer[:, None, :2] + xy + tot[:, :, 2:] = outer[:, None, 2:] + inner[None, :, 2:] # sum rotations and mirrored + tot[:, :, 2] %= 2 * pi # clamp rot + tot[:, :, 3] %= 2 # clamp mirrored + + if tensor: + return tot + return numpy.concatenate(tot)