diff --git a/masque/pattern.py b/masque/pattern.py index f120c52..7b39d97 100644 --- a/masque/pattern.py +++ b/masque/pattern.py @@ -15,12 +15,15 @@ from .subpattern import SubPattern from .repetition import GridRepetition from .shapes import Shape, Polygon from .label import Label -from .utils import rotation_matrix_2d, vector2 +from .utils import rotation_matrix_2d, vector2, normalize_mirror from .error import PatternError __author__ = 'Jan Petykiewicz' +visitor_function_t = Callable[['Pattern', Tuple['Pattern'], Dict, numpy.ndarray], 'Pattern'] + + class Pattern: """ 2D layout consisting of some set of shapes and references to other Pattern objects @@ -165,6 +168,78 @@ class Pattern: pat = memo[pat_id] return pat + def dfs(self, + visit_before: visitor_function_t = None, + visit_after: visitor_function_t = None, + transform: numpy.ndarray or bool or None = False , + memo: Dict = None, + hierarchy: Tuple['Pattern'] = (), + ) -> 'Pattern': + """ + Experimental convenience function. + Performs a depth-first traversal of this pattern and its subpatterns. + At each pattern in the tree, the following sequence is called: + ``` + current_pattern = visit_before(current_pattern, **vist_args) + for sp in current_pattern.subpatterns] + sp.pattern = sp.pattern.df(visit_before, visit_after, updated_transform, + memo, (current_pattern,) + hierarchy) + current_pattern = visit_after(current_pattern, **visit_args) + ``` + where `visit_args` are + `hierarchy`: (top_pattern, L1_pattern, L2_pattern, ..., parent_pattern) + tuple of all parent-and-higher patterns + `transform`: numpy.ndarray containing cumulative + [x_offset, y_offset, rotation (rad), mirror_x (0 or 1)] + for the instance being visited + `memo`: Arbitrary dict (not altered except by visit_*()) + + :param visit_before: Function to call before traversing subpatterns. + Should accept a Pattern and **visit_args, and return the (possibly modified) + pattern. Default None (not called). + :param visit_after: Function to call after traversing subpatterns. + Should accept a Pattern and **visit_args, and return the (possibly modified) + pattern. Default None (not called). + :param transform: Initial value for `visit_args['transform']`. + Can be `False`, in which case the transform is not calculated. + `True` or `None` is interpreted as [0, 0, 0, 0]. + :param memo: Arbitrary dict for use by visit_*() functions. Default None (empty dict). + :param hierarchy: Tuple of patterns specifying the hierarchy above the current pattern. + Appended to the start of the generated `visit_args['hierarchy']`. + Default is an empty tuple. + """ + if memo is None: + memo = {} + + if transform is None or transform is True: + transform = numpy.zeros(4) + + if self in hierarchy: + raise PatternError('.dfs() called on pattern with circular reference') + + pat = self + if visit_before is not None: + pat = visit_before(pat, hierarchy=hierarchy, memo=memo, transform=transform) + + for subpattern in self.subpatterns: + if transform is not False: + mirror_x, angle = normalize_mirror(subpattern.mirrored) + angle += subpattern.rotation + sp_transform = transform + numpy.hstack((subpattern.offset, angle, mirror_x)) + sp_transform[3] %= 2 + else: + sp_transform = False + + subpattern.pattern = subpattern.pattern.dfs(visit_before=visit_before, + visit_after=visit_after, + transform=sp_transform, + memo=memo, + hierarchy=hierarchy + (self,)) + + if visit_after is not None: + pat = visit_after(pat, hierarchy=hierarchy, memo=memo, transform=transform) + return pat + def polygonize(self, poly_num_points: int = None, poly_max_arclen: float = None, diff --git a/masque/utils.py b/masque/utils.py index bf82cba..aa99bf2 100644 --- a/masque/utils.py +++ b/masque/utils.py @@ -57,6 +57,20 @@ def rotation_matrix_2d(theta: float) -> numpy.ndarray: [numpy.sin(theta), +numpy.cos(theta)]]) +def normalize_mirror(mirrored: Tuple[bool, bool]) -> Tuple[bool, float]: + mirrored_x, mirrored_y = mirrored + if mirrored_x and mirrored_y: + angle = numpy.pi + mirror_x = False + elif mirrored_x: + angle = 0 + mirror_x = True + elif mirror_y: + angle = numpy.pi + mirror_x = True + return mirror_x, angle + + def remove_duplicate_vertices(vertices: numpy.ndarray, closed_path: bool = True) -> numpy.ndarray: duplicates = (vertices == numpy.roll(vertices, 1, axis=0)).all(axis=1) if not closed_path: