diff --git a/docs/programs.md b/docs/programs.md index 8dd61cc2..90dbc797 100644 --- a/docs/programs.md +++ b/docs/programs.md @@ -41,6 +41,8 @@ Arbitrary tools for viewing, analyzing, or working with instamatic data. - [instamatic.defocus_helper](#instamaticdefocus_helper) (`instamatic.gui.defocus_button:main`) - [instamatic.find_crystals](#instamaticfind_crystals) (`instamatic.processing.find_crystals:main_entry`) - [instamatic.find_crystals_ilastik](#instamaticfind_crystals_ilastik) (`instamatic.processing.find_crystals_ilastik:main_entry`) +- [instamatic.grid_finder](#instamaticgridfinder) (`instamatic.grid.finder:main`) +- [instamatic.grid_monitor](#instamaticgridmonitor) (`instamatic.grid.monitor:main`) - [instamatic.learn](#instamaticlearn) (`scripts.learn:main_entry`) **Server** @@ -546,6 +548,55 @@ d `, `--output `, `--output_name ` : Generate `MapScaleInd.yaml` for `predicrystal` from config. +## instamatic.grid_finder + +Determine the geometry and orientation of the copper grid using various methods. + +**Usage:** +```bash +instamatic.grid_finder [-h] [FILEPATH] manual +instamatic.grid_finder [-h] [FILEPATH] auto [--idx IDX] [--arms ARMS] + [--order ORDER] [--offset OFFSET] +``` +**Positional arguments:** + +`FILEPATH` +: Path to the file where grid geometry is or is to be stored, typically `grid.yaml`. + +**Optional arguments:** + +`-h`, `--help` +: Show this help message and exit + +`--idx` +: Target window index to find (default: next available window) + +`--arms` +: Number of directions the sweep arms will go to find edge + +`--order` +: For each order above 1, sweep also in previous orders midpoints + +`--offset` +: Rotate the arms by this many degrees before first order search + +`--video` +: Display the live video stream during automated grid finding + + +## instamatic.grid_monitor + +GUI program to monitor a grid.yaml file and plot any updates live. + +**Usage:** +```bash +instamatic.grid_monitor [-h] [FILEPATH] +``` +**Positional arguments:** + +`FILEPATH` +: Path to the file where grid geometry is stored, typically `grid.yaml`. + ## instamatic.learn Predict whether a crystal is of good or bad quality by its diffraction pattern. diff --git a/pyproject.toml b/pyproject.toml index 722996cf..6f124fbe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,8 @@ publishing = [ "instamatic.defocus_helper" = "instamatic.gui.defocus_button:main" "instamatic.find_crystals" = "instamatic.processing.find_crystals:main_entry" "instamatic.find_crystals_ilastik" = "instamatic.processing.find_crystals_ilastik:main_entry" +"instamatic.grid_finder" = "instamatic.grid.finder:main" +"instamatic.grid_monitor" = "instamatic.grid.monitor:main" "instamatic.learn" = "scripts.learn:main_entry" # server "instamatic.temserver" = "instamatic.server.tem_server:main" diff --git a/src/instamatic/camera/camera_serval.py b/src/instamatic/camera/camera_serval.py index 9922428d..3b8ef2f0 100644 --- a/src/instamatic/camera/camera_serval.py +++ b/src/instamatic/camera/camera_serval.py @@ -81,7 +81,7 @@ def establish_connection(self) -> ServalCamera: conn.set_detector_config(**self.detector_config) conn.destination = {'Image': [http_dest]} - if getattr(self, 'stream_movies_via_tcp') or STREAM_MOVIES_VIA_TCP: + if getattr(self, 'stream_movies_via_tcp', None) or STREAM_MOVIES_VIA_TCP: self.tcp_listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.tcp_listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self.tcp_listener.settimeout(1.0) diff --git a/src/instamatic/grid/__init__.py b/src/instamatic/grid/__init__.py new file mode 100644 index 00000000..7ffe73e4 --- /dev/null +++ b/src/instamatic/grid/__init__.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import Optional, Union + +import numpy as np + +Intercepts = dict[int, np.ndarray] + + +def cross2d(a: np.ndarray, b: np.ndarray) -> float: + """A scalar 2d cross product between two arrays of length 2.""" + return (a[0] * b[1] - a[1] * b[0]).item() + + +def versor( + *, + deg: Optional[Union[float, np.ndarray]] = None, + rad: Optional[Union[float, np.ndarray]] = None, +) -> np.ndarray: + """A versor in the direction of angle expressed in radians or degrees.""" + radians = np.deg2rad(deg) if rad is None else rad + return np.array([np.cos(radians), np.sin(radians)], dtype=float) diff --git a/src/instamatic/grid/artist.py b/src/instamatic/grid/artist.py new file mode 100644 index 00000000..25eebce3 --- /dev/null +++ b/src/instamatic/grid/artist.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from typing import Optional + +import numpy as np +from matplotlib import pyplot as plt +from matplotlib.axes import Axes +from matplotlib.figure import Figure +from matplotlib.patches import Polygon +from matplotlib.ticker import FuncFormatter + +from instamatic._typing import float_nm +from instamatic.grid import Intercepts +from instamatic.grid.grid import PeriodicConvexPolygonGrid + + +def plot_grid( + grid: PeriodicConvexPolygonGrid, + *, + intercepts: Optional[Intercepts] = None, + limit_x: Optional[float_nm] = None, + limit_y: Optional[float_nm] = None, + ax: Optional[Axes] = None, + show_indices: bool = True, + show_intercepts: bool = True, + figsize: tuple[float, float] = (5, 5), + dpi: int = 100, +) -> tuple[Figure, Axes]: + """Draw geometry with windows based on intercepts dict or limit_x/y.""" + + fig, ax = (ax.figure, ax) if ax else plt.subplots(figsize=figsize, dpi=dpi) + + fig.patch.set_facecolor('black') + ax.set_facecolor('black') + ax.set_aspect('equal', adjustable='box') + ax.tick_params(colors='white', direction='out') + for spine in ax.spines.values(): + spine.set_color('white') + ax.set_xlabel('x / um', color='white') + ax.set_ylabel('y / um', color='white') + ax.xaxis.set_major_formatter(FuncFormatter(lambda x, _: f'{x * 1e-3:g}')) + ax.yaxis.set_major_formatter(FuncFormatter(lambda y, _: f'{y * 1e-3:g}')) + ax.grid(True, which='major', color='white', linewidth=0.8, alpha=0.25, zorder=0) + ax.set_axisbelow(True) + patch_kw = {'facecolor': 'white', 'edgecolor': 'white', 'closed': True, 'zorder': 1} + text_kw = {'color': 'black', 'ha': 'center', 'va': 'center', 'fontsize': 10, 'zorder': 2} + + indices = sorted(intercepts) if intercepts else list(range(25)) + if limit_x is not None and limit_y is not None: + indices = grid.windows_in_limits(x=limit_x, y=limit_y) + + cmap = plt.colormaps.get_cmap('tab10') + for idx in indices: + try: + window = grid.window(idx) + except ValueError: # negative/invalid Ulam index + color = '#808080' + else: + color = cmap(idx % 10) + corners = np.asarray(window.corners, dtype=float) + ax.add_patch(Polygon(corners, **patch_kw)) + if show_indices: + cx, cy = map(float, window.center) + ax.text(cx, cy, str(idx), **text_kw) + + if show_intercepts and intercepts and idx in intercepts: + xys = np.asarray(intercepts[idx], dtype=float) + ax.plot( + xys[:, 0], + xys[:, 1], + marker='x', + linestyle='', + color=color, + markersize=6, + zorder=5, + ) + + ax.relim() + ax.autoscale_view() + + x0, x1 = ax.get_xlim() + y0, y1 = ax.get_ylim() + cx, cy = 0.5 * (x0 + x1), 0.5 * (y0 + y1) + r = 0.5 * max(x1 - x0, y1 - y0) + ax.set_xlim(cx - r, cx + r) + ax.set_ylim(cy - r, cy + r) + + ax.set_autoscale_on(False) # Freeze limits so lines don't affect view + ax.axhline(0, color='white', linewidth=1.0, alpha=0.6, zorder=0) + ax.axvline(0, color='white', linewidth=1.0, alpha=0.6, zorder=0) + + if limit_x is not None: + ax.axvline(-limit_x, color='red', linewidth=1.0, zorder=4) + ax.axvline(limit_x, color='red', linewidth=1.0, zorder=4) + if limit_y is not None: + ax.axhline(-limit_y, color='red', linewidth=1.0, zorder=4) + ax.axhline(limit_y, color='red', linewidth=1.0, zorder=4) + + return fig, ax diff --git a/src/instamatic/grid/finder.py b/src/instamatic/grid/finder.py new file mode 100644 index 00000000..963ea4f3 --- /dev/null +++ b/src/instamatic/grid/finder.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +import argparse +from math import sqrt +from pathlib import Path +from textwrap import dedent +from typing import TYPE_CHECKING, Optional + +import numpy as np +import yaml + +from instamatic._typing import AnyPath, float_nm, int_nm +from instamatic.grid import Intercepts +from instamatic.grid.grid import GRID_REGISTRY, PeriodicConvexPolygonGrid + +if TYPE_CHECKING: + from instamatic.gui.click_dispatcher import ClickListener + + +class GridFinder: + """Base strategy for determining and updating grid geometry. + Can be written to or read from a yaml file in the following format: + + grid_type: square + geometry: {x: 11111, y: 22222, t: 33.333, w: 44444, h: null, s: 55555} + intercepts: + 0: [[1001, 2002], [3003, 4004], ] # ... + 1: [[5005, 6006], [7007, 8008], ] # ... + -1: [[88005, 88006], [99007, 99008], ] + # negative window number = intercepts not assigned to window yet ... + """ + + GRID_REGISTRY_INV = {v: k for k, v in GRID_REGISTRY.items()} + + def __init__( + self, + grid: Optional[PeriodicConvexPolygonGrid] = None, + intercepts: Optional[Intercepts] = None, + ) -> None: + self.grid = grid or GRID_REGISTRY['square'](0, 0, 0, 85_000) + self.intercepts: Intercepts = intercepts or {} + self.path: Optional[AnyPath] = None # if present, auto-save here + + @classmethod + def from_yaml(cls, yaml_path: AnyPath) -> GridFinder: + with open(Path(yaml_path), 'r') as f: + data = yaml.safe_load(f) + grid = GRID_REGISTRY[data['grid_type']](**data['geometry']) + intercepts = data.get('intercepts', {}) + return cls(grid, {k: np.array(v, dtype=float) for k, v in intercepts.items()}) + + def to_yaml(self, yaml_path: AnyPath) -> None: + grid_type_name = self.GRID_REGISTRY_INV[type(self.grid)] + data = { + 'grid_type': grid_type_name, + 'geometry': self.grid.to_params(), + 'intercepts': {k: v.tolist() for k, v in self.intercepts.items()}, + } + with open(yaml_path, 'w') as f: + yaml.dump(data, f, default_flow_style=None, sort_keys=False) + + def add_intercept(self, window_idx: int, x: float_nm, y: float_nm) -> None: + """Register a new intercept of given id, x, and y in the finder.""" + if window_idx in self.intercepts: + self.intercepts[window_idx] = np.vstack([self.intercepts[window_idx], [x, y]]) + else: + self.intercepts[window_idx] = np.array([[x, y]], dtype=float) + if self.path is not None: + self.to_yaml(self.path) + + def fit_intercepts(self, window_idx: int) -> None: + """Fit all intercepts with given window id to a new window.""" + try: + xy = self.intercepts[window_idx] + except KeyError as e: + raise KeyError(f'No intercepts for {window_idx=} found!') from e + if 0 in self.intercepts: + new_center = (np.max(xy, axis=0) + np.min(xy, axis=0)) / 2 + new_window_idx = self.grid.nearest_index(*new_center) + else: + new_window_idx = 0 + self.grid = type(self.grid).guess({0: xy}) + new_intercepts = self.intercepts[window_idx] + del self.intercepts[window_idx] + if new_window_idx in self.intercepts: + oi = self.intercepts[window_idx] + self.intercepts[new_window_idx] = np.vstack([oi, new_intercepts]) + else: + self.intercepts[new_window_idx] = new_intercepts + self.grid.refine(self.intercepts) + if self.path is not None: + self.to_yaml(self.path) + + def refine_by_manual_clicking(self, ctrl, cl: ClickListener) -> None: + """Update grid & intercepts via clicks when stage is at window edge. + + Move the stage to as many points on one windows edge as possible + (at least the corners and midpoints). At each point, click LMB + to add current stage position as one edge point. RMB to finish. + """ + from instamatic.gui.click_dispatcher import MouseButton + + print(dedent(self.refine_by_manual_clicking.__doc__)) + while True: + prev_grid, prev_intercepts = self.grid, self.intercepts + with cl: + while True: + c = cl.get_click() + if c.button == MouseButton.RIGHT: + break + self.add_intercept(-1, *ctrl.stage.xy) + + self.fit_intercepts(-1) + print('Intercepts fit: LMB to accept, RMB to retry, MMB for new window') + c = cl.get_click() + if c.button == MouseButton.LEFT: + break + elif c.button == MouseButton.RIGHT: + self.grid, self.intercepts = prev_grid, prev_intercepts + if self.path is not None: + self.to_yaml(self.path) + + def refine_by_auto_sweeping( + self, + ctrl, + window_idx: int = -1, + x_lim: Optional[int_nm] = None, + y_lim: Optional[int_nm] = None, + arms: int = 3, + order: Optional[int] = None, + offset: Optional[float] = None, + ) -> None: + """Let grid & intercepts refine by automatically looking for edges. + + Move to `window_idx` or next window (if any present, else start here). + If the requested window is predicted to lie inside a bounding box span + by `x_lim` and `y_lim`, look for the edges by monitoring total beam + intensity. `arms`, `order`, `offset` determine `star_sweep` precision. + """ + from instamatic.grid.sweeping import star_sweep + + idx = window_idx + if not self.intercepts: + idx = 0 + else: + d_lim = (sqrt(abs(max(self.intercepts))) + 2) * (self.grid.w + self.grid.h) + x_lim = x_lim or d_lim # crude estimate of new window search area + y_lim = y_lim or d_lim # if no limits was given: (sqrt(idx)+2)(w+h) + idc_in_limits = self.grid.windows_in_limits(x=x_lim, y=y_lim) + if idx == -1: + try: + idx = min([i for i in idc_in_limits if i not in self.intercepts]) + except ValueError: + raise IndexError('Could not locate next window within limits') + else: + if idx not in idc_in_limits: + raise IndexError(f'Requested window {idx} is not within limits') + + if idx > 0: + ctrl.stage.set(*[int(xy) for xy in self.grid.window(idx).center]) + + smart_order = 3 if idx == 0 else 2 if len(self.intercepts.keys()) < 4 else 1 + ss_order = smart_order if order is None else order + ss_offset = offset if offset is not None else 17 * idx + + for xy in star_sweep(arms=arms, order=ss_order, offset=ss_offset): + self.add_intercept(idx, *xy) + self.fit_intercepts(idx) + + +def main(): + """CLI tool to determine grid geometry using various methods.""" + + from instamatic.controller import initialize + + parser = argparse.ArgumentParser(description=main.__doc__) + parser.add_argument( + '-f', + '--file', + type=str, + help='A custom path to the grid.yaml file with results', + default='grid.yaml', + ) + + subparsers = parser.add_subparsers( + dest='method', + required=True, + help='Method of grid geometry determination', + ) + + _ = subparsers.add_parser('manual', help='Manual via moving stage & input') + a = subparsers.add_parser('auto', help='Automatically via star sweeping') + + a.add_argument( + '--idx', + type=int, + default=-1, + help='Target window index to find (default: next available window)', + ) + + a.add_argument( + '--arms', + type=int, + default=3, + choices=[3, 4, 5, 6, 7], + help='Number of directions the sweep arms will go to find edge', + ) + + a.add_argument( + '--order', + type=int, + default=None, + choices=[1, 2, 3, 4, 5], + help='For each order above 1, sweep also in previous orders midpoints', + ) + + a.add_argument( + '--offset', + type=float, + default=None, + help='Rotate the arms by this many degrees before first order search', + ) + + a.add_argument( + '--video', + action='store_true', + help='Display the live video stream during automated grid finding', + ) + + args = parser.parse_args() + + try: + gf = GridFinder.from_yaml(args.file) + except FileNotFoundError: + gf = GridFinder() + gf.path = args.file + ctrl = initialize() + + def run_grid_finder(_cl: Optional[ClickListener] = None) -> None: + """Run the finder logic, without blocking the GUI if it is needed.""" + try: + if args.method == 'manual': + gf.refine_by_manual_clicking(ctrl, _cl) + elif args.method == 'auto': + gf.refine_by_auto_sweeping( + ctrl=ctrl, + window_idx=args.idx, + arms=args.arms, + order=args.order, + offset=args.offset, + ) + finally: + if use_video: + root.after(0, root.destroy) + + use_video = args.method == 'manual' or (args.method == 'auto' and args.video) + if use_video: + from threading import Thread + from tkinter import Tk + + from instamatic import config + from instamatic.camera import LiveVideoStream + from instamatic.gui.videostream_frame import VideoStreamFrame + + root = Tk() + stream = LiveVideoStream(cam=config.camera.name) + vsf = VideoStreamFrame(root, stream=stream) + vsf.pack(side='top', fill='both', expand=True) + cl = vsf.click_dispatcher.add_listener(name='grid_finder', active=True) + Thread(target=run_grid_finder, daemon=True, args=(cl,)).start() + root.mainloop() + else: + run_grid_finder() + + +if __name__ == '__main__': + main() diff --git a/src/instamatic/grid/grid.py b/src/instamatic/grid/grid.py new file mode 100644 index 00000000..6a1528b9 --- /dev/null +++ b/src/instamatic/grid/grid.py @@ -0,0 +1,309 @@ +from __future__ import annotations + +from typing import Annotated, Generic, Optional, Protocol, TypeVar, Union + +import numpy as np +from scipy.optimize import least_squares +from typing_extensions import Self + +from instamatic._collections import NoOverwriteDict +from instamatic._typing import float_deg, float_nm +from instamatic.grid import versor +from instamatic.grid.window import ( + GridablePolygonWindow, + HexagonalWindow, + RectangularWindow, + SquareWindow, +) +from instamatic.utils.pairing import hulam2uv, ij2ulam, ulam2ij, uv2hulam + +DualIndex = tuple[int, int] +SingleIndex = Annotated[int, 'positive'] +WindowIndex = Union[DualIndex, SingleIndex] +WindowType = TypeVar('WindowType', bound=GridablePolygonWindow) + + +class PairingFunction(Protocol): + def __call__(self, i: int, j: int, /) -> int: ... + + +class PairingInverse(Protocol): + def __call__(self, n: int, /) -> tuple[int, int]: ... + + +WindowGeometryTuple = tuple[float_nm, float_nm, float_deg, float_nm, Optional[float_nm]] + + +class PeriodicConvexPolygonGrid(Generic[WindowType]): + """A ConvexPolygonGrid with identical windows and on a 2D ab-lattice. + + The conventional, most-expected lattice kind for ED experiments. + Every window is an identical convex polygon placed in the same + distance from other windows, as determined by the grid support + thickness. Utilizes internal coordinate system of its "central" + window 0, with two axes, "a" & "b", selected in such a way that the + angle between axes "a" and coordinate X is minimal, and the angle + from "a" to "b" is positive (clockwise) and minimal. The length of + "a" and "b" should match expected distance to next windows. + """ + + neighborhood: np.ndarray[int] + pairing_function: PairingFunction + pairing_inverse: PairingInverse + window_type: type[WindowType] + + DEFAULT_SPACING: float_nm = 10_000 # underestimated on purpose (real ~35k) + + def __init__( + self, + x: float_nm, # x coordinate of the grid origin in stage coordinates + y: float_nm, # y coordinate of the grid origin in stage coordinates + t: float_deg, # signed angle from the X-axis to a-vector in degrees + w: float_nm, # length of X-aligned axis: edge to edge center-points + h: Optional[float_nm] = None, # length of the other axis, if relevant + s: Optional[float_nm] = None, # spacing between neighbor grid windows + ): + self.x = x + self.y = y + self.t = t + self.w = w + self._h = h + self._s = s + + def __repr__(self) -> str: + params = ', '.join(f'{k}: {v}' for k, v in self.to_params().items()) + return f'{self.__class__.__name__}({params})' + + def normalized(self) -> Self: + """Align w with X axis by casting theta to [+,- interior angle / 2]""" + a = float(self.window_type.INTERIOR_ANGLE) + n = int(np.floor((self.t + 0.5 * a) / a)) # rotates needed to min theta + t = float(self.t - n * a) + w = abs(float(self.w)) + h = None if self._h is None else abs(float(self._h)) + if self._h is not None and (n % 2): + w, h = h, w + return self.__class__(self.x, self.y, t, w, h, self._s) + + @property + def origin(self) -> np.ndarray: + """Origin of the grid i.e. its window 0 in stage coordinates (nm).""" + return np.array([self.x, self.y], dtype=float) + + @property + def h(self): + """Value of "h" if it is applicable or "w" in square/hex cases.""" + return self._h if self.window_type.USES_HEIGHT else self.w + + @h.setter + def h(self, value: float_nm) -> None: + self._h = value if self.window_type.USES_HEIGHT else None + + @property + def s(self): + """Uniform spacing between two neighbor windows i.e. grid thickness.""" + return self.DEFAULT_SPACING if self._s is None else self._s + + @s.setter + def s(self, value: float_nm) -> None: + self._s = value + + @property + def a_dir(self) -> np.ndarray: + """A versor oriented along grid space axis "a" in stage coords.""" + t = float(np.radians(self.t)) + return np.array([np.cos(t), np.sin(t)], dtype=float) + + @property + def a_edge(self) -> np.ndarray: + """Half-window vector from center to edge midpoint along axis "a".""" + return (self.w / 2) * self.a_dir + + @property + def a_grid(self) -> np.ndarray: + """Center-to-center lattice vector to neighboring window along "a".""" + return (self.w + self.s) * self.a_dir + + @property + def b_dir(self) -> np.ndarray: + """A versor oriented along grid space axis "b" in stage coords.""" + t = float(np.radians(self.t + self.window_type.INTERIOR_ANGLE)) + return np.array([np.cos(t), np.sin(t)], dtype=float) + + @property + def b_edge(self) -> np.ndarray: + """Half-window vector from center to edge midpoint along axis "b".""" + return (self.h / 2) * self.b_dir + + @property + def b_grid(self) -> np.ndarray: + """Center-to-center lattice vector to neighboring window along "b".""" + return (self.h + self.s) * self.b_dir + + def window_geometry(self, idx: WindowIndex) -> WindowGeometryTuple: + """Return the current geom: origin + shape params of window "idx".""" + ij: DualIndex = idx if isinstance(idx, tuple) else self.pairing_inverse(idx) + x, y = self.origin + ij[0] * self.a_grid + ij[1] * self.b_grid + return x, y, self.t, self.w, self.h + + def window(self, idx: WindowIndex) -> WindowType: + """Convenience method that makes a window located at requested idx.""" + return self.window_type(*self.window_geometry(idx)) + + def windows_in_limits(self, x: float_nm, y: float_nm) -> list[int]: + """List indices of windows intersecting the box [-x, x] x [-y, y].""" + candidates_idx: set[int] = {0, self.nearest_index(0.0, 0.0)} + idx_in_limits: list[int] = [] + + while candidates_idx: + idx = min(candidates_idx) + candidates_idx.remove(idx) + + if self.window(idx=idx).intersects_limits(x, y): + idx_in_limits.append(idx) + for nb in np.array(self.pairing_inverse(idx)) + self.neighborhood: + nb_idx = self.pairing_function(int(nb[0]), int(nb[1])) + if nb_idx > idx: + candidates_idx.add(nb_idx) + + return idx_in_limits + + def nearest_index(self, x: float_nm, y: float_nm) -> int: + """Return spiral index of predicted window nearest to the center.""" + delta = np.asarray([x, y], dtype=float) - self.origin + metric = np.column_stack([self.a_grid, self.b_grid]) + ij, *_ = np.linalg.lstsq(metric, delta, rcond=None) + i, j = (int(np.rint(v)) for v in ij) + return int(self.pairing_function(i, j)) + + @classmethod + def guess(cls, intercepts: dict[int, np.ndarray]) -> Self: + """Guess the geometry of window 0 given points on its edges.""" + if 0 not in intercepts: + raise ValueError('No intercepts for window 0 provided') + + xys0 = np.asarray(intercepts[0], dtype=float) + c0 = np.mean(xys0, axis=0) + deltas = xys0 - c0 + half_span = 0.5 * cls.window_type.INTERIOR_ANGLE + thetas = np.linspace(-half_span, +half_span, 91) + + best_guess = None + best_score = np.inf + + for t in thetas: + a_dir = versor(deg=t) + b_dir = versor(deg=t + cls.window_type.INTERIOR_ANGLE) + qa = deltas @ a_dir + qb = deltas @ b_dir + + if cls.window_type.USES_HEIGHT: + w = 2.0 * np.quantile(np.abs(qa), 0.9) + h = 2.0 * np.quantile(np.abs(qb), 0.9) + else: + w = 2.0 * np.quantile(np.hstack([np.abs(qa), np.abs(qb)]), 0.9) + h = None + + g = cls(x=c0[0], y=c0[1], t=t, w=w, h=h, s=None) + score = float(np.sum(g.window(0).edge_residuals(xys0) ** 2)) + if score < best_score: + best_guess = g + best_score = score + + assert best_guess is not None + return best_guess + + def refine(self, intercepts: dict[int, np.ndarray]) -> None: + """Refine self to match the window_id: intercepts dictionary.""" + + windows = sorted(intercepts.keys()) + refine_spacing = len(windows) > 1 + fit_h = self.window_type.USES_HEIGHT + fixed_s = self._s # preserve "unknown spacing" when not refined + + def serialize(g: PeriodicConvexPolygonGrid) -> np.ndarray: + """Express the geometry instance as a series of refined vars.""" + vals = [float(g.x), float(g.y), float(g.t), float(g.w)] + if fit_h: + vals.append(float(g.h)) + if refine_spacing: + vals.append(float(g.s)) + return np.asarray(vals, dtype=float) + + def deserialize(p: np.ndarray) -> PeriodicConvexPolygonGrid: + """Convert a series of refined vars into a periodic geometry.""" + vals = iter(p) + x = float(next(vals)) + y = float(next(vals)) + t = float(next(vals)) + w = float(next(vals)) + h = float(next(vals)) if fit_h else None + s = float(next(vals)) if refine_spacing else fixed_s + return self.__class__(x=x, y=y, t=t, w=w, h=h, s=s).normalized() + + def residuals(p: np.ndarray) -> np.ndarray: + """Calculate residual for each window in refined geometry.""" + geom = deserialize(p) + res: list[np.ndarray] = [] + for idx, xys in intercepts.items(): + tmp_window = geom.window(idx) + res.append(tmp_window.edge_residuals(xys)) + return np.concatenate(res) if res else np.empty(0, dtype=float) + + lower = [-np.inf, -np.inf, -np.inf, 1e-9] + upper = [np.inf, np.inf, np.inf, np.inf] + if fit_h: + lower.append(1e-9) + upper.append(np.inf) + if refine_spacing: + lower.append(0.0) + upper.append(np.inf) + res = least_squares( + residuals, + x0=serialize(self), + bounds=(np.asarray(lower, dtype=float), np.asarray(upper, dtype=float)), + method='trf', + loss='soft_l1', + f_scale=10_000, + ) + + geometry = deserialize(res.x) + if not refine_spacing: + geometry._s = fixed_s # keep spacing unknown/fixed in the 1-window case + + self.x = geometry.x + self.y = geometry.y + self.t = geometry.t + self.w = geometry.w + self.h = geometry.h + self.s = geometry._s + + def to_params(self): + return {'x': self.x, 'y': self.y, 't': self.t, 'w': self.w, 'h': self._h, 's': self._s} + + +class HexagonalGrid(PeriodicConvexPolygonGrid): + neighborhood = np.array([(1, 0), (0, 1), (-1, 1), (-1, 0), (0, -1), (1, -1)], dtype=int) + pairing_function: PairingFunction = staticmethod(uv2hulam) + pairing_inverse: PairingInverse = staticmethod(hulam2uv) + window_type: type[WindowType] = HexagonalWindow + + +class RectangularGrid(PeriodicConvexPolygonGrid[RectangularWindow]): + neighborhood = np.array([(1, 0), (0, 1), (-1, 0), (0, -1)], dtype=int) + pairing_function: PairingFunction = staticmethod(ij2ulam) + pairing_inverse: PairingInverse = staticmethod(ulam2ij) + window_type: type[WindowType] = RectangularWindow + + +class SquareGrid(PeriodicConvexPolygonGrid[SquareWindow]): + neighborhood = np.array([(1, 0), (0, 1), (-1, 0), (0, -1)], dtype=int) + pairing_function: PairingFunction = staticmethod(ij2ulam) + pairing_inverse: PairingInverse = staticmethod(ulam2ij) + window_type: type[WindowType] = SquareWindow + + +GRID_REGISTRY = NoOverwriteDict[str, type[PeriodicConvexPolygonGrid]]() +GRID_REGISTRY['hexagonal'] = HexagonalGrid +GRID_REGISTRY['rectangular'] = RectangularGrid +GRID_REGISTRY['square'] = SquareGrid diff --git a/src/instamatic/grid/monitor.py b/src/instamatic/grid/monitor.py new file mode 100644 index 00000000..476098dd --- /dev/null +++ b/src/instamatic/grid/monitor.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import argparse +import tkinter as tk +from pathlib import Path +from tkinter import filedialog, ttk +from typing import Optional + +import matplotlib.pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + +from instamatic._typing import AnyPath +from instamatic.grid.artist import plot_grid +from instamatic.grid.finder import GridFinder + + +class GridMonitor(ttk.Frame): + """A Tkinter-based GUI to monitor and refine TEM grids.""" + + def __init__(self, parent: tk.Widget, grid_path: Optional[AnyPath] = None) -> None: + super().__init__(parent) + self.parent = parent + self.path = tk.StringVar(value=str(grid_path) if grid_path else '') + self.last_mtime: float = 0.0 + + canvas_frame = ttk.Frame(self) + canvas_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + self.fig, self.ax = plt.subplots(figsize=(5, 5), dpi=100) + self.canvas = FigureCanvasTkAgg(self.fig, master=canvas_frame) + self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True) + + address_frame = ttk.Frame(self) + address_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=5, pady=5) + ttk.Label(address_frame, text='Grid File:').pack(side=tk.LEFT) + + path_entry = ttk.Entry(address_frame, textvariable=self.path) + path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5) + browse_btn = ttk.Button(address_frame, text='Browse', command=self._browse_file) + browse_btn.pack(side=tk.LEFT) + + self.pack(fill=tk.BOTH, expand=True) + self.after(1000, self._poll_file_updates) + + def _browse_file(self): + filepath = filedialog.askopenfilename( + title='Select YAML with grid information', + filetypes=(('YAML files', '*.yaml *.yml'), ('All files', '*.*')), + ) + if filepath: + self.path.set(filepath) + self.last_mtime = 0.0 # Force a redraw + + def _poll_file_updates(self): + """Continuously check if the grid YAML file has been modified.""" + grid_path = Path(self.path.get()) + if grid_path and grid_path.exists(): + try: + current_mtime = grid_path.stat().st_mtime + if current_mtime > self.last_mtime: + self.last_mtime = current_mtime + self._redraw_grid() + except OSError: + pass + + self.after(1000, self._poll_file_updates) + + def _redraw_grid(self): + """Load the latest YAML and redraw the matplotlib canvas.""" + grid_path = Path(self.path.get()) + if not grid_path or not grid_path.exists(): + return + + self.fig.clf() + self.ax = self.fig.add_subplot(111) + gf = GridFinder.from_yaml(grid_path) + plot_grid( + grid=gf.grid, + intercepts=gf.intercepts, + ax=self.ax, + show_indices=True, + show_intercepts=True, + ) + self.canvas.draw_idle() + + +def main(): + """GUI program to monitor a grid.yaml file and plot any updates live.""" + parser = argparse.ArgumentParser(description=main.__doc__) + parser.add_argument( + '-f', '--file', type=str, help='Path to the grid.yaml file', default=None + ) + + args = parser.parse_args() + + root = tk.Tk() + root.title('Instamatic Grid Monitor') + _ = GridMonitor(root, grid_path=args.file) + root.mainloop() + + +if __name__ == '__main__': + main() diff --git a/src/instamatic/grid/sweeping.py b/src/instamatic/grid/sweeping.py new file mode 100644 index 00000000..4e9ae6ab --- /dev/null +++ b/src/instamatic/grid/sweeping.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from dataclasses import dataclass +from itertools import chain +from typing import Any, Generator, Literal, Optional, Sequence, Union + +import numpy as np +from typing_extensions import Self + +from instamatic._typing import float_deg, float_nm, int_nm +from instamatic.controller import TEMController, _ctrl, initialize +from instamatic.grid import cross2d, versor +from instamatic.utils.iterating import pairwise + +if not _ctrl: + _ctrl: TEMController = initialize() + + +Vector2 = Sequence[float] + + +class InstanceAutoNameRegistry: + """Autosave each subclass instance in `cls.INSTANCES` dict under `name`""" + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + cls.INSTANCES: dict[Any, Self] = {} + + def __post_init__(self): + self.__class__.INSTANCES[getattr(self, 'name')] = self + + +@dataclass +class SweeperTeam(InstanceAutoNameRegistry): + """Stores a set of shared variables between the members of sweeper team.""" + + name: str = '' # identifier used for registration in INSTANCES + step_size: int_nm = 10_000 # largest step size allowed + precision: int_nm = 1 # smallest step size allowed + threshold: float = 0.1 # fraction of light_max that signals the edge + light_max: int = -1 # maximum light observed at any point by any sweeper + sampling: int = 2 # confirm edge only if this many consecutive steps are dark + + +default_sweeper_team = SweeperTeam() + + +class Sweeper: + """A simple descriptor of stage movement with fixed heading.""" + + def __init__(self, origin: Vector2, heading: Vector2, team: str = '') -> None: + self.origin = np.array([origin[0], origin[1]], dtype=float) + self.heading = np.array([heading[0], heading[1]], dtype=float) + self.position = np.array([origin[0], origin[1]], dtype=float) + self.team = SweeperTeam.INSTANCES[team] + + def breed(self, other: Self) -> Self: + """Return a new instance with mean heading and position.""" + o = (self.position + other.position) / 2 + n = np.linalg.norm(s := self.heading + other.heading) + if n == 0: + raise ValueError('Cannot breed sweepers with parallel heading') + return self.__class__(origin=o, heading=s / n, team=self.team.name) + + def dist2segment(self, x1: float, y1: float, x2: float, y2: float) -> float: + """Dist to intercept a segment or +inf, stackoverflow.com/q/2931573.""" + ray1o = self.origin + ray1h = self.heading + ray2o = np.array([x1, y1], dtype=float) + ray2h = np.array([x2 - x1, y2 - y1], dtype=float) + delta = ray2o - ray1o + cross = cross2d(ray1h, ray2h) + if cross == 0: + return np.inf + ray1d = cross2d(delta, ray2h) / cross + ray2d = cross2d(delta, ray1h) / cross + if ray1d >= 0 and 0 <= ray2d <= 1: + return ray1d * np.linalg.norm(self.heading) + return np.inf + + def peak(self) -> int: + """Return light (image sum) at current position, update light max.""" + light = int(_ctrl.get_image(header_keys=())[0].sum(dtype=np.int64)) + self.team.light_max = max(light, self.team.light_max) + return light + + def goto(self, x: int_nm, y: int_nm) -> None: + """Change sweeper position to `x`, `y` and update current position.""" + _ctrl.stage.set(x=x, y=y) + self.position = np.array([x, y], dtype=float) + + def step(self, length: float_nm) -> None: + """Change sweeper position by `length` in `heading` direction.""" + x0, y0 = _ctrl.stage.xy + x1 = int(x0 + self.heading[0].item() * length) + y1 = int(y0 + self.heading[1].item() * length) + self.goto(x1, y1) + + +class BinarySweeper(Sweeper): + """A stage-state descriptor used to binary-search the grid edge.""" + + def sweep(self) -> None: + """Bin-search the edge based on peaked light vs max * threshold.""" + self.goto(x=int(self.origin[0]), y=int(self.origin[1])) + step_size: float_nm = self.team.step_size + samples: int = 0 + while True: # linear search - step forward until `samples` dark frames + if self.peak() < self.team.threshold * self.team.light_max: + samples += 1 + else: + samples = 0 + if samples >= self.team.sampling: + break + self.step(step_size) + step_size *= 0.5 + self.step(length=-(2 * samples - 1) * step_size) # half step past 1st dark + while step_size > self.team.precision: # binary search + step_size *= 0.5 + is_lit = self.peak() > self.team.threshold * self.team.light_max + self.step(length=step_size if is_lit else -step_size) + + +def star_sweep( + arms: Literal[3, 4, 5, 6, 7] = 3, + order: Literal[1, 2, 3, 4, 5] = 3, + offset: float_deg = 0, +) -> Generator[Vector2, None, None]: + """Send sweepers in a star shape, yield xy on edge as they are found. + + arms: Number of unique directions to send initial "order=1" sweepers in. + order: For each above 1, send new sweepers in bisectors of previous order. + offset: Applied when defining "order=1" headings to add more variety. + + With default settings, 1 sweeper takes ~30 seconds, less for higher order. + The total number of sweepers = arms * 2 ** (order - 1): 12 with defaults. + """ + center: Vector2 = np.array(_ctrl.stage.xy, dtype=int) + team = str(center) + _ = SweeperTeam(name=team) + + # define headings, sweep with initial sweepers to approx grid center (order=1) + headings = offset + np.linspace(0, 360, num=arms, endpoint=False, dtype=float) + done_sweepers = [] + for h in headings: + bs = BinarySweeper(origin=center, heading=versor(deg=h), team=team) + bs.sweep() + yield bs.position + done_sweepers.append(bs) + + # For each order above 1, generate bisectors of previous orders, sweep, yield + for _ in range(1, order): + new_sweepers = [] + for a, b in pairwise(done_sweepers, closed=True): + ns = a.breed(b) + ns.sweep() + yield ns.position + new_sweepers.append(ns) + done_sweepers = list(chain.from_iterable(zip(done_sweepers, new_sweepers))) diff --git a/src/instamatic/grid/window.py b/src/instamatic/grid/window.py new file mode 100644 index 00000000..d1481ada --- /dev/null +++ b/src/instamatic/grid/window.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Optional + +import numpy as np +from typing_extensions import Self + +from instamatic._typing import float_deg, float_nm +from instamatic.grid import versor +from instamatic.utils.iterating import pairwise + + +class Window(ABC): + """Describes an arbitrary single window on a TEM grid.""" + + center: np.ndarray = ... # 2-element array describing the center of window + + +class ConvexPolygonWindow(Window): + """Describes any convex polygon TEM grid window with known corners.""" + + corners: np.ndarray = ... # a Nx2 ordered array of xy corner coordinates + + def x_intersections(self, y: float_nm) -> Optional[tuple[float, float]]: + """Return (x_min, x_max) for a horizontal line intersecting at y.""" + intersection_xs: list[float] = [] + for (x1, y1), (x2, y2) in pairwise(self.corners, closed=True): + if y1 == y2: # edge case (degeneracy / double counting) + continue + intersection_fraction = (y - y1) / (y2 - y1) + if not 0 <= intersection_fraction < 1: + continue # does not intersect + intersection_xs.append(x1 + (x2 - x1) * intersection_fraction) + if len(intersection_xs) < 2: + return None + return min(intersection_xs), max(intersection_xs) + + def y_intersections(self, x: float_nm) -> Optional[tuple[float, float]]: + """Return (y_min, y_max) for a vertical line intersecting at x.""" + intersection_ys: list[float] = [] + for (x1, y1), (x2, y2) in pairwise(self.corners, closed=True): + if x1 == x2: # edge case (degeneracy / double counting) + continue + intersection_fraction = (x - x1) / (x2 - x1) + if not 0 <= intersection_fraction < 1: + continue # does not intersect + intersection_ys.append(y1 + (y2 - y1) * intersection_fraction) + if len(intersection_ys) < 2: + return None + return min(intersection_ys), max(intersection_ys) + + +class GridablePolygonWindow(ConvexPolygonWindow): + """Describes a polygon window with a 2D (a, b) grid coordinate system. + + This kind of window is expected to exist in a periodic grid, + therefore it should include an internal coordinate system with two + axes "a" and "b". They should be selected in such a way that the + angle between axes "a" and X is minimal and the angle from "a" to + "b" is positive and minimal. The length of "a" and "b" should match + the distance between window center and its edge. + + Any subclass of GridablePolygonWindow should initialize using at least + four following parameters in this order, and other as needed: + + - x: x coordinate of the window center in the stage XY coordinate system; + - y: y coordinate of the window center in the stage XY coordinate system; + - t: smallest signed angle from stage +X axis towards "a" axis in degrees; + - w: double the distance between window center and its' edge midpoint; + """ + + INTERIOR_ANGLE: float_deg = ... # class attribute: angle between a and b axes + USES_HEIGHT: bool = ... # True if a secondary metric i.e. height is needed + a: np.ndarray = ... + b: np.ndarray = ... # from center towards the side, not aligned with ~X + + def __init__( + self, + x: float_nm, + y: float_nm, + t: float_deg, + w: float_nm, + h: Optional[float_nm] = None, + ) -> None: + """A uniform abstract constructor for all subclasses (nm/degrees).""" + self.x: float_nm = float(x) + self.y: float_nm = float(y) + self.t: float_deg = float(t) + self.w: float_nm = float(w) + self.h: float_nm = self.w if h is None else float(h) + + self.a: np.ndarray = ... # vector aligned with ~X direction + self.b: np.ndarray = ... # "a" rotated by INTERIOR_ANGLE anti-clockwise + self.corners: np.ndarray = ... # ordered anti-clockwise, start from "a" + + def __repr__(self) -> str: + """Accurate representation, show params as floats (from to_params).""" + parts = [f'{k}={float(v)}' for k, v in self.to_params().items()] + return f'{type(self).__name__}(' + ', '.join(parts) + ')' + + @property + def center(self) -> np.ndarray: + return np.array([self.x, self.y], dtype=float) + + def edge_residuals(self, xys: np.ndarray) -> np.ndarray: + """Return residual distance to the nearest edge per point.""" + xys = np.asarray(xys, dtype=float) + if len(xys) == 0: + return np.empty(0, dtype=float) + + p1 = np.asarray(self.corners, dtype=float) # (M, 2) + edge_vecs = np.roll(p1, -1, axis=0) - p1 # (M, 2) + edge_l2 = np.sum(edge_vecs * edge_vecs, axis=1) # (M,) + + if np.any(edge_l2 == 0): + raise ValueError('Degenerate polygon edge: consecutive corners coincide') + + # Vector from each segment start to each point + rel = xys[:, None, :] - p1[None, :, :] # (N, M, 2) + + # Projection parameter onto each edge, then clamp to the finite segment + t = np.sum(rel * edge_vecs[None, :, :], axis=2) / edge_l2[None, :] # (N, M) + t = np.clip(t, 0.0, 1.0) + + # Closest point on each segment + closest = p1[None, :, :] + t[:, :, None] * edge_vecs[None, :, :] # (N, M, 2) + + # Distance from each point to each segment + dists = np.linalg.norm(xys[:, None, :] - closest, axis=2) # (N, M) + + return np.min(dists, axis=1) + + def intersects_limits(self, x: float_nm, y: float_nm) -> bool: + """Test whether the window intersects the box [-x, x] x [-y, y]. To + this aim, in seven consecutive blocks: + + 1) Alias the corners and their coordinates for further + convenience; 2) Test whether the window bounding box (min/max) + is beyond limits; 3) Test whether any window corner is inside + the limits; 4) to 7) Test if any limit line intersects the + window within limits. + """ + c = np.asarray(self.corners, dtype=float) # window corners + cx = c[:, 0] # view of window corners' x coordinates + cy = c[:, 1] # view of window corners' x coordinates + + if cx.max() < -x or cx.min() > x or cy.max() < -y or cy.min() > y: + return False + + if np.any((cx > -x) & (cx < x) & (cy > -y) & (cy < y)): + return True + + xs = self.x_intersections(y=y) + if xs is not None and xs[0] < x and xs[1] > -x: + return True + + xs = self.x_intersections(y=-y) + if xs is not None and xs[0] < x and xs[1] > -x: + return True + + ys = self.y_intersections(x=x) + if ys is not None and ys[0] < y and ys[1] > -y: + return True + + ys = self.y_intersections(x=-x) + if ys is not None and ys[0] < y and ys[1] > -y: + return True + + return False + + @abstractmethod + def to_params(self) -> dict[str, float]: ... + + +class HexagonalWindow(GridablePolygonWindow): + """A regular hexagonal window with a 2D "ab" coordinate system.""" + + INTERIOR_ANGLE: float_deg = 60.0 + ROT60MAT = np.array([[1, -np.sqrt(3)], [np.sqrt(3), 1]], dtype=float) / 2 + USES_HEIGHT = False + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + self.a = 0.5 * self.w * versor(deg=self.t) + self.b = self.ROT60MAT @ self.a + + angles = self.t + np.array([30, 90, 150, 210, 270, 330], dtype=float) + self.corners = self.center + self.w / np.sqrt(3.0) * versor(deg=angles).T + + def to_params(self) -> dict[str, float]: + return {'x': self.x, 'y': self.y, 't': self.t, 'w': self.w} + + def translated(self, delta: np.ndarray) -> Self: + """Return a new window translated by (dx, dy) in nm.""" + d = np.asarray(delta, dtype=float) + return type(self)(self.x + d[0], self.y + d[1], self.t, self.w) + + +class RectangularWindow(GridablePolygonWindow): + """Describes one rectangular window with a 2D ab coordinate system. + + Geometry is described using five immutable float scalars (nm / degree): + + - center_x: coordinate of the window center on the X axis; + - center_y: coordinate of the window center on the Y axis; + - width: length of window side aligned with the direction of X axis; + - height: length of window side aligned with the direction or Y axis; + - theta: signed angle from X axis towards A axis and the X-aligned edge. + """ + + INTERIOR_ANGLE: float_deg = 90.0 + USES_HEIGHT = True + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + c = self.center + self.a = a = 0.5 * self.w * versor(deg=self.t) + self.b = b = 0.5 * self.h * versor(deg=self.t + 90) + self.corners = np.vstack([c + a + b, c - a + b, c - a - b, c + a - b]) + + def to_params(self) -> dict[str, float]: + return {'x': self.x, 'y': self.y, 't': self.t, 'w': self.w, 'h': self.h} + + def translated(self, delta: np.ndarray) -> Self: + """Return a new window translated by (dx, dy) in nm.""" + d = np.asarray(delta, dtype=float) + return type(self)(self.x + d[0], self.y + d[1], self.t, self.w, self.h) + + +class SquareWindow(RectangularWindow): + """A regular square window with a 2D "ab" coordinate system.""" + + USES_HEIGHT = False + + def to_params(self) -> dict[str, float]: + return {'x': self.x, 'y': self.y, 't': self.t, 'w': self.w} + + def translated(self, delta: np.ndarray) -> Self: + """Return a new window translated by (dx, dy) in nm.""" + d = np.asarray(delta, dtype=float) + return type(self)(self.x + d[0], self.y + d[1], self.t, self.w) diff --git a/tests/test_grid.py b/tests/test_grid.py new file mode 100644 index 00000000..5bbc782c --- /dev/null +++ b/tests/test_grid.py @@ -0,0 +1,251 @@ +from __future__ import annotations + +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import matplotlib.pyplot as plt +import numpy as np +import pytest +from matplotlib.axes import Axes +from matplotlib.figure import Figure + +from instamatic.grid import cross2d, versor +from instamatic.grid.artist import plot_grid +from instamatic.grid.finder import GridFinder +from instamatic.grid.grid import ( + HexagonalGrid, + PeriodicConvexPolygonGrid, + RectangularGrid, + SquareGrid, +) +from instamatic.grid.sweeping import Sweeper +from instamatic.grid.window import ( + GridablePolygonWindow, + HexagonalWindow, + RectangularWindow, + SquareWindow, +) +from tests.utils import InstanceAutoTracker + +# instamatic.grid.__init__ + + +def test_grid_cross2d(): + """Assert that cross2d product is calculated correctly.""" + a = np.array([1.0, 0.0]) + a2 = np.array([1.0, 0.0]) + b = np.array([0.0, 1.0]) + z = np.array([0.0, 0.0]) + assert cross2d(a, b) == pytest.approx(1.0) + assert cross2d(a, a2) == pytest.approx(0.0) + assert cross2d(a, z) == pytest.approx(0.0) + + +def test_grid_versors(): + """Assert that versors are created properly and reasonably.""" + np.testing.assert_allclose(versor(rad=0), [1, 0], atol=1e-12) + np.testing.assert_allclose(versor(deg=90), [0, 1], atol=1e-12) + for angle in np.linspace(0, 360, 37): + assert np.linalg.norm(versor(deg=angle)) == pytest.approx(1.0) + v1 = versor(deg=30.0) + v2 = versor(rad=0.5235987756) + v3 = np.array([np.sqrt(3) / 2, 1 / 2], dtype=np.float32) + np.testing.assert_allclose(v1, v3, atol=1e-6) + np.testing.assert_allclose(v2, v3, atol=1e-6) + + +# instamatic.grid.window + + +@dataclass +class WindowTestCase(InstanceAutoTracker): + """Auto-registers three windows test case instances in INSTANCES.""" + + cls: type[GridablePolygonWindow] + params: tuple[float, ...] + h_cut: float = 10 # length of horizontal / vertical line cutting through middle + v_cut: float = 10 + + +r2 = np.sqrt(2) +r3 = np.sqrt(3) + +WindowTestCase(cls=HexagonalWindow, params=(30, 30, 0, 10), h_cut=20 / r3) +WindowTestCase(cls=RectangularWindow, params=(30, 30, 0, 10, 10)) +WindowTestCase(cls=SquareWindow, params=(30, 30, 45, 10), h_cut=10 * r2, v_cut=10 * r2) + + +@pytest.mark.parametrize('window_case', WindowTestCase.INSTANCES) +def test_windows(window_case) -> None: + """Test if created window meets simple tests for internal lengths.""" + w = window_case.cls(*window_case.params) + assert w.corners.shape[0] > 3 + np.testing.assert_allclose(w.center, [30, 30]) + side_lens = np.linalg.norm(np.roll(w.corners, -1, axis=0) - w.corners, axis=1) + dists = np.linalg.norm(w.corners - w.center, axis=1) + if window_case.cls == HexagonalWindow: + np.testing.assert_allclose(side_lens, 10 / r3, rtol=1e-10) + np.testing.assert_allclose(dists, 10 / r3, rtol=1e-10) + else: + np.testing.assert_allclose(side_lens, 10, rtol=1e-10) + np.testing.assert_allclose(dists, 5 * r2, rtol=1e-10) + + +@pytest.mark.parametrize('window_case', WindowTestCase.INSTANCES) +def test_window_residuals(window_case) -> None: + """Test if edge residual calculations work and give 0 for corners.""" + w = window_case.cls(*window_case.params) + np.testing.assert_allclose(w.edge_residuals(w.corners), 0.0, atol=1e-6) + assert w.edge_residuals(np.empty((0, 2))).shape == (0,) + + +@pytest.mark.parametrize('window_case', WindowTestCase.INSTANCES) +def test_window_intersects(window_case) -> None: + """Test if windows are correctly cut or missed by intersecting lines.""" + w = window_case.cls(*window_case.params) + wx = w.y_intersections(30) + assert wx[1] - wx[0] == pytest.approx(window_case.h_cut) + wy = w.x_intersections(30) + assert wy[1] - wy[0] == pytest.approx(window_case.v_cut) + assert w.x_intersections(-30) is None + assert w.y_intersections(-30) is None + + +@pytest.mark.parametrize('window_case', WindowTestCase.INSTANCES) +def test_window_hexagonal_intersects_limits(window_case) -> None: + """Test if windows intersect correct x/y limits' squares.""" + w = window_case.cls(*window_case.params) + assert not w.intersects_limits(5, 5) + assert w.intersects_limits(30, 30) + + +# instamatic.grid.grid + + +@dataclass +class GridTestCase(InstanceAutoTracker): + """Auto-registers grid test case instances in INSTANCES.""" + + cls: type[PeriodicConvexPolygonGrid] + window_type: type + h: Optional[float] = None + s: Optional[float] = None + + +GridTestCase(cls=HexagonalGrid, window_type=HexagonalWindow, h=None, s=5.0) +GridTestCase(cls=RectangularGrid, window_type=RectangularWindow, h=50) +GridTestCase(cls=SquareGrid, window_type=SquareWindow, h=None, s=None) + + +@pytest.mark.parametrize('grid_case', GridTestCase.INSTANCES) +def test_grid_initialization_and_normalization(grid_case) -> None: + """Test that grids initialize correctly and normalization bounds the + angle.""" + g = grid_case.cls(10, 20, 30, 40, grid_case.h, grid_case.s) + if grid_case.s is None: + assert g.s == g.DEFAULT_SPACING + g_norm = g.normalized() + max_angle = g.window_type.INTERIOR_ANGLE / 2 + assert -max_angle <= g_norm.t <= max_angle + + +@pytest.mark.parametrize('grid_case', GridTestCase.INSTANCES) +def test_grid_window_and_nearest_index(grid_case) -> None: + """Test whether initialized window type and nearest index are correct.""" + g = grid_case.cls(10, 20, 30, 40, grid_case.h, grid_case.s) + w = g.window(77) + assert isinstance(w, grid_case.window_type) + assert g.nearest_index(*w.center) == 77 + + +@pytest.mark.parametrize('grid_case', GridTestCase.INSTANCES) +def test_grid_limits(grid_case) -> None: + """Test that limit bounding boxes properly encapsulate targeted windows.""" + g = grid_case.cls(10, 20, 30, 40, grid_case.h, grid_case.s) + idc_in_limits = g.windows_in_limits(20, 20) + assert 0 in idc_in_limits + assert 77 not in idc_in_limits + + +@pytest.mark.parametrize('grid_case', GridTestCase.INSTANCES) +def test_grid_guess_and_refine(grid_case) -> None: + """Test that guessing and refining a grid from perfect intercepts recovers + it.""" + g = grid_case.cls(0, 10, 20, 30, grid_case.h, grid_case.s) + + corners = g.window(0).corners # window corners + midpoints = (corners + np.roll(corners, -1, axis=0)) / 2 + intercepts = {0: np.vstack((corners, midpoints))} + g_guess = grid_case.cls.guess(intercepts) + np.testing.assert_allclose(g_guess.origin, g.origin, atol=1e-7) + g_guess.refine(intercepts) + for a in ('origin', 'x', 'y', 't', 'w', 'h'): + np.testing.assert_allclose(getattr(g_guess, a), getattr(g, a), atol=1e-10) + + +# instamatic.grid.finder + + +def test_grid_finder_add_intercept(): + """Assert that new intercepts are properly saved in GridFinder.""" + gf = GridFinder(grid=SquareGrid(0, 0, 0, 10_000), intercepts={}) + gf.add_intercept(0, 200.0, 0) + assert 0 in gf.intercepts + assert gf.intercepts[0].shape == (1, 2) + gf.add_intercept(0, 400.0, 0) + gf.add_intercept(0, 600.0, 0) + assert gf.intercepts[0].shape == (3, 2) + + +def test_grid_finder_yaml_read_write(): + """Assert that grid finder can read and (auto-)write yaml files.""" + gf = GridFinder(grid=SquareGrid(0, 0, 0, 10_000), intercepts={}) + gf.add_intercept(0, 1000.0, 2000.0) + with tempfile.TemporaryDirectory() as d: + p = Path(d) / 'grid.yaml' + gf.to_yaml(p) + gf2 = GridFinder.from_yaml(p) + assert type(gf2.grid) is SquareGrid + np.testing.assert_allclose(gf2.grid.x, gf.grid.x) + np.testing.assert_allclose(gf2.intercepts[0], gf.intercepts[0]) + gf.path = p + gf.add_intercept(1, 5000.0, 5000.0) + gf3 = GridFinder.from_yaml(p) + assert type(gf3.grid) is SquareGrid + np.testing.assert_allclose(gf2.grid.w, gf.grid.w) + np.testing.assert_allclose(gf3.intercepts[1], gf.intercepts[1]) + + +# instamatic.grid.sweeping + + +def test_grid_sweeping_sweeper_dist2segment(): + """Assert dist2segment properly estimates dist only to hit segments.""" + s = Sweeper(origin=[0, 0], heading=[1, 0]) + assert s.dist2segment(4, -1, 6, 1) == pytest.approx(5.0) + assert s.dist2segment(0, 1, 10, 1) == np.inf + assert s.dist2segment(-5, -1, -5, 1) == np.inf + + +# instamatic.grid.artist + + +def test_grid_plot_grid(): + """Assert that plot_grid executes with default arguments and returns + Figure/Axes.""" + fig, ax = plot_grid( + SquareGrid(0, 0, 0, 10_000), + intercepts={0: np.atleast_2d([0.0, 0.0])}, + limit_x=25_000, + limit_y=25_000, + ) + assert isinstance(fig, Figure) + assert isinstance(ax, Axes) + assert len(ax.patches) > 0 # Should have drawn the default windows 0-24 + assert len(ax.lines) >= 6 + plt.close(fig) + + +# instamatic.grid.monitor is GUI only and thus has no tests