From 1ab3bd75838f9be8fd6d61679f6081ec1e33e4f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Tue, 1 Sep 2026 20:13:46 +0200 Subject: [PATCH 1/7] Add grid from side branch, add docs and entry points --- docs/programs.md | 48 ++++ pyproject.toml | 2 + src/instamatic/grid/__init__.py | 22 ++ src/instamatic/grid/artist.py | 99 +++++++++ src/instamatic/grid/finder.py | 270 +++++++++++++++++++++++ src/instamatic/grid/grid.py | 380 ++++++++++++++++++++++++++++++++ src/instamatic/grid/monitor.py | 102 +++++++++ src/instamatic/grid/sweeping.py | 169 ++++++++++++++ src/instamatic/grid/window.py | 257 +++++++++++++++++++++ tests/test_grid.py | 251 +++++++++++++++++++++ 10 files changed, 1600 insertions(+) create mode 100644 src/instamatic/grid/__init__.py create mode 100644 src/instamatic/grid/artist.py create mode 100644 src/instamatic/grid/finder.py create mode 100644 src/instamatic/grid/grid.py create mode 100644 src/instamatic/grid/monitor.py create mode 100644 src/instamatic/grid/sweeping.py create mode 100644 src/instamatic/grid/window.py create mode 100644 tests/test_grid.py diff --git a/docs/programs.md b/docs/programs.md index 8dd61cc2..88a22552 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,52 @@ 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 + + +## 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/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..06f2aa8b --- /dev/null +++ b/src/instamatic/grid/finder.py @@ -0,0 +1,270 @@ +from __future__ import annotations + +import argparse +from math import sqrt +from pathlib import Path +from textwrap import dedent +from typing import 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 +from instamatic.gui.click_dispatcher import ClickEvent, ClickListener, MouseButton + + +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 # nm + t: 0.033 # degrees + w: 44444 + s: 6666 + intercepts: + 0: + - [1001, 2002] + - [3003, 4004] + # ... + -1: # negative number = fresh entry, not assigned to a window yet + - [88005, 88006] + - [99007, 99008] + # ... + # ... + """ + + 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, 50_000, 50_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.""" + xy = self.intercepts[window_idx] + 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. + + Navigate the stage to as many points on one windows edge as + possible (at least the corners and midpoints). At each point, + position the edge at the center of the screen and LMB to add the + point. RMB to finish. + """ + 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: Optional[int] = None, + order: Optional[int] = None, + offset: Optional[float] = None, + ) -> None: + """Let grid & intercepts refine by automatically looking for edges. + + Navigate to `window_idx` or next window and, if it is inside a bounding + box span by `x_lim` and `y_lim`, look for the edges by monitoring total + 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') + + 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_arms = arms if arms is not None else 3 + ss_offset = offset if offset is not None else 17 * idx + + for xy in star_sweep(arms=ss_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=3, + 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=0.0, + help='Rotate the arms by this many degrees before first order search', + ) + + args = parser.parse_args() + + # Initialize finding logic and controller + gf = GridFinder.from_yaml(args.file) + gf.path = args.file + ctrl = initialize() + + if args.method == 'manual': + + class TerminalClickListener: + """Mocks ClickListener for CLI by mapping keyboard inputs to + MouseButtons.""" + + def __enter__(self) -> TerminalClickListener: + print('Entering terminal mock for mouse button click listener.') + print('Press "ENTER" for LMB, R+ENTER for RMB, M+ENTER for MMB.') + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + pass + + def get_click(self) -> ClickEvent: + """Prompt user for terminal input to simulate mouse clicks.""" + from instamatic.gui.click_dispatcher import ClickEvent, MouseButton + + cmd = input('>> ').strip().lower() + + if cmd == 'r': + return ClickEvent(button=MouseButton.RIGHT) + elif cmd == 'm': + return ClickEvent(button=MouseButton.MIDDLE) + return ClickEvent(button=MouseButton.LEFT) + + cl = TerminalClickListener() + gf.refine_by_manual_clicking(ctrl, cl) + + elif args.command == 'auto': + gf.refine_by_auto_sweeping( + ctrl=ctrl, window_idx=args.idx, arms=args.arms, order=args.order, offset=args.offset + ) + + +if __name__ == '__main__': + main() diff --git a/src/instamatic/grid/grid.py b/src/instamatic/grid/grid.py new file mode 100644 index 00000000..59ffa88f --- /dev/null +++ b/src/instamatic/grid/grid.py @@ -0,0 +1,380 @@ +from __future__ import annotations + +from typing import Annotated, Generic, Optional, Protocol, Self, TypeVar, Union + +import numpy as np +from scipy.optimize import least_squares + +from instamatic._collections import NoOverwriteDict +from instamatic._typing import float_nm +from instamatic.grid.window import ( + GridablePolygonWindow, + HexagonalWindow, + RectangularWindow, + SquareWindow, +) +from instamatic.utils.pairing import hulam2uv, ij2ulam, ulam2ij, uv2hulam + +DualIndex = tuple[int, int] +SpiralIndex = Annotated[int, 'positive'] +WindowIndex = Union[DualIndex, SpiralIndex] +WindowType = TypeVar('WindowType', bound=GridablePolygonWindow) + + +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) + + +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, float_nm, Optional[float_nm]] +WindowShapeTuple = tuple[float, 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 + + 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, # signed angle from the X-axis towards 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 + + +if __name__ == '__main__': + import numpy as np + + residuals = { + 0: np.array( + [ + [572543.92272, 458611.73068], + [564971.6688, 458611.73068], + [564971.6688, 458611.73068], + [564971.6688, 458611.73068], + [552876.06528, 458609.70234], + [552876.06528, 458609.70234], + [538110.39744, 458608.68817], + [538110.39744, 458608.68817], + [538110.39744, 458608.68817], + [518853.20256, 458610.71651], + [518853.20256, 458610.71651], + [496719.8544, 458613.75902], + [496719.8544, 458613.75902], + [487814.08368, 434983.59802], + [484349.97072, 421938.32931], + [484349.97072, 421938.32931], + [483886.27056, 411164.8014], + [483886.27056, 411164.8014], + [486744.23952, 397380.20276], + [486744.23952, 397380.20276], + [486744.23952, 397380.20276], + [486183.55632, 380419.22368], + [486183.55632, 380419.22368], + [496566.80304, 365588.0016], + [496566.80304, 365588.0016], + [503120.73504, 365593.07245], + [503120.73504, 365593.07245], + [511500.67584, 365799.96313], + [511500.67584, 365799.96313], + [520806.5016, 366044.3781], + [520806.5016, 366044.3781], + [532121.69472, 366439.9044], + [543770.26704, 366439.9044], + [543770.26704, 366439.9044], + [555862.83984, 366764.4388], + [555862.83984, 366764.4388], + [568560.04128, 367445.96104], + [568560.04128, 367445.96104], + [572369.65632, 406690.28336], + [576968.77392, 419188.91444], + [575013.95952, 431882.26616], + [572419.6632, 441778.53702], + [572422.69392, 454949.56281], + ] + ) + } + a = SquareGrid(10000, 20000, 0, 50000) + print(a.guess(residuals).to_params()) + print(a.to_params()) + a.refine(residuals) + print(a.to_params()) + a = a.guess(residuals) + a.refine(residuals) + print(a.to_params()) 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..370d6551 --- /dev/null +++ b/src/instamatic/grid/sweeping.py @@ -0,0 +1,169 @@ +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 + + +class Sweeper: + """A simple descriptor of stage movement with fixed heading.""" + + def __init__(self, origin: Vector2, heading: Vector2) -> 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) + + 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 + + +@dataclass +class EdgeSweeperTeam(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 + + +default_edge_sweeper_team = EdgeSweeperTeam() + + +class EdgeSweeper(Sweeper): + """Used to determine the edge of the stage based on camera feedback.""" + + def __init__(self, origin: Vector2, heading: Vector2, team: str = '') -> None: + self.history: list[Vector2] = [] + self.team = EdgeSweeperTeam.INSTANCES[team] + super().__init__(origin, heading) + + 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.history.append([x, 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 MarchingEdgeSweeper(EdgeSweeper): + """Moves monotonously, stops when intensity fract < max * threshold.""" + + def sweep(self) -> None: + """Walk steps into heading until peaked light is below threshold.""" + self.goto(x=int(self.origin[0]), y=int(self.origin[1])) + light_here: int = self.peak() + while light_here > self.team.light_max * self.team.threshold: + self.step(length=self.team.step_size) + light_here = self.peak() + + +class BinaryEdgeSweeper(EdgeSweeper): + """A stage-state descriptor used to binary-search the grid edge.""" + + 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 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 + direction: Literal[1, -1] + refining: bool = False + while step_size > self.team.precision: + light_here = self.peak() + if light_here > self.team.threshold * self.team.light_max: + direction = 1 + else: + direction = -1 + refining = True + if refining: + step_size *= 0.5 + self.step(length=direction * step_size) + + +def star_sweep( + arms: Literal[3, 4, 5, 6, 7] = 5, + order: Literal[1, 2, 3, 4, 5] = 5, + offset: float_deg = 0, +) -> Generator[Vector2, None, None]: + """Sweep window, yielding (x, y) points on its edge as they are found.""" + center: Vector2 = np.array(_ctrl.stage.xy, dtype=int) + team = str(center) + _ = EdgeSweeperTeam(name=team) + + # define and sweep with initial sweepers to approximate grid center (order=1) + headings = offset + np.linspace(0, 360, num=arms, endpoint=False, dtype=float) + directions = [versor(deg=h) for h in headings] + bes = [BinaryEdgeSweeper(origin=center, heading=d, team=team) for d in directions] + + for sweeper in bes: + sweeper.sweep() + yield sweeper.position + + # For each order above 1, generate bisectors of previous orders, sweep, yield + finished_bes = bes + for _ in range(1, order): + new_bes = [] + for a, b in pairwise(finished_bes, closed=True): + ns = a.breed(b) + ns.sweep() + yield ns.position + new_bes.append(ns) + finished_bes = list(chain.from_iterable(zip(finished_bes, new_bes))) diff --git a/src/instamatic/grid/window.py b/src/instamatic/grid/window.py new file mode 100644 index 00000000..f4b0b00c --- /dev/null +++ b/src/instamatic/grid/window.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Optional, Union + +import numpy as np +from typing_extensions import Self + +from instamatic._typing import float_nm +from instamatic.grid import versor +from instamatic.utils.iterating import pairwise + +X = np.array([1, 0], dtype=float) +Y = np.array([0, 1], dtype=float) + + +WindowGeometryTuple = tuple[float_nm, float_nm, float, float_nm, Optional[float_nm]] + + +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 = ... # 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, + 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 = 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).""" + p = self.to_params() + parts = [f'{k}={float(v)}' for k, v in p.items()] + return f'{type(self).__name__}(' + ', '.join(parts) + ')' + + def __str__(self) -> str: + """Nicely display self, show params as integers (from to_params).""" + p = self.to_params() + parts = [f'{k}={int(np.rint(float(v)))}' for k, v in p.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 = 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 = 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 From 8da704a7f52faab11b02b76336674506b3ef654f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Tue, 1 Sep 2026 20:33:39 +0200 Subject: [PATCH 2/7] Remove redundant versor definition --- src/instamatic/grid/grid.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/instamatic/grid/grid.py b/src/instamatic/grid/grid.py index 59ffa88f..68585d0e 100644 --- a/src/instamatic/grid/grid.py +++ b/src/instamatic/grid/grid.py @@ -7,6 +7,7 @@ from instamatic._collections import NoOverwriteDict from instamatic._typing import float_nm +from instamatic.grid import versor from instamatic.grid.window import ( GridablePolygonWindow, HexagonalWindow, @@ -21,16 +22,6 @@ WindowType = TypeVar('WindowType', bound=GridablePolygonWindow) -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) - - class PairingFunction(Protocol): def __call__(self, i: int, j: int, /) -> int: ... From 42854014d4c6f42e5e18546377186247fb2e3554 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Tue, 1 Sep 2026 20:51:32 +0200 Subject: [PATCH 3/7] Delay imports in finder --- src/instamatic/grid/finder.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/instamatic/grid/finder.py b/src/instamatic/grid/finder.py index 06f2aa8b..4f3f0bc5 100644 --- a/src/instamatic/grid/finder.py +++ b/src/instamatic/grid/finder.py @@ -4,7 +4,7 @@ from math import sqrt from pathlib import Path from textwrap import dedent -from typing import Optional +from typing import TYPE_CHECKING, Optional import numpy as np import yaml @@ -12,7 +12,9 @@ from instamatic._typing import AnyPath, float_nm, int_nm from instamatic.grid import Intercepts from instamatic.grid.grid import GRID_REGISTRY, PeriodicConvexPolygonGrid -from instamatic.gui.click_dispatcher import ClickEvent, ClickListener, MouseButton + +if TYPE_CHECKING: + from instamatic.gui.click_dispatcher import ClickListener class GridFinder: @@ -104,6 +106,8 @@ def refine_by_manual_clicking(self, ctrl, cl: ClickListener) -> None: position the edge at the center of the screen and LMB to add the 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 @@ -232,6 +236,7 @@ def main(): ctrl = initialize() if args.method == 'manual': + from instamatic.gui.click_dispatcher import ClickEvent, MouseButton class TerminalClickListener: """Mocks ClickListener for CLI by mapping keyboard inputs to @@ -247,8 +252,6 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None: def get_click(self) -> ClickEvent: """Prompt user for terminal input to simulate mouse clicks.""" - from instamatic.gui.click_dispatcher import ClickEvent, MouseButton - cmd = input('>> ').strip().lower() if cmd == 'r': From 06630c285646a28ad5e34016fa6cce18a29a3c51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Wed, 2 Sep 2026 12:04:33 +0200 Subject: [PATCH 4/7] Import Self from typing_extensions to avoid Python3.9/10 issues --- src/instamatic/grid/grid.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/instamatic/grid/grid.py b/src/instamatic/grid/grid.py index 68585d0e..70db07aa 100644 --- a/src/instamatic/grid/grid.py +++ b/src/instamatic/grid/grid.py @@ -1,9 +1,10 @@ from __future__ import annotations -from typing import Annotated, Generic, Optional, Protocol, Self, TypeVar, Union +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_nm From 05fab9d6c0c022e7abbca5bbc10fb16fe80d7b62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Wed, 2 Sep 2026 15:30:28 +0200 Subject: [PATCH 5/7] Remove unused `MarchingEdgeSweeper` --- src/instamatic/grid/sweeping.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/instamatic/grid/sweeping.py b/src/instamatic/grid/sweeping.py index 370d6551..60b93051 100644 --- a/src/instamatic/grid/sweeping.py +++ b/src/instamatic/grid/sweeping.py @@ -97,18 +97,6 @@ def step(self, length: float_nm) -> None: self.goto(x1, y1) -class MarchingEdgeSweeper(EdgeSweeper): - """Moves monotonously, stops when intensity fract < max * threshold.""" - - def sweep(self) -> None: - """Walk steps into heading until peaked light is below threshold.""" - self.goto(x=int(self.origin[0]), y=int(self.origin[1])) - light_here: int = self.peak() - while light_here > self.team.light_max * self.team.threshold: - self.step(length=self.team.step_size) - light_here = self.peak() - - class BinaryEdgeSweeper(EdgeSweeper): """A stage-state descriptor used to binary-search the grid edge.""" From b25099dbdce8ca890344c4571d799a7b21c9f34d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Wed, 2 Sep 2026 15:57:19 +0200 Subject: [PATCH 6/7] Clear, optimize `sweeping.py` code, add documentation --- src/instamatic/grid/sweeping.py | 98 ++++++++++++++++----------------- 1 file changed, 48 insertions(+), 50 deletions(-) diff --git a/src/instamatic/grid/sweeping.py b/src/instamatic/grid/sweeping.py index 60b93051..0aa8e437 100644 --- a/src/instamatic/grid/sweeping.py +++ b/src/instamatic/grid/sweeping.py @@ -30,13 +30,36 @@ 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 + + +default_sweeper_team = SweeperTeam() + + class Sweeper: """A simple descriptor of stage movement with fixed heading.""" - def __init__(self, origin: Vector2, heading: Vector2) -> None: + 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.""" @@ -54,29 +77,6 @@ def dist2segment(self, x1: float, y1: float, x2: float, y2: float) -> float: return ray1d * np.linalg.norm(self.heading) return np.inf - -@dataclass -class EdgeSweeperTeam(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 - - -default_edge_sweeper_team = EdgeSweeperTeam() - - -class EdgeSweeper(Sweeper): - """Used to determine the edge of the stage based on camera feedback.""" - - def __init__(self, origin: Vector2, heading: Vector2, team: str = '') -> None: - self.history: list[Vector2] = [] - self.team = EdgeSweeperTeam.INSTANCES[team] - super().__init__(origin, heading) - 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)) @@ -86,7 +86,6 @@ def peak(self) -> int: 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.history.append([x, y]) self.position = np.array([x, y], dtype=float) def step(self, length: float_nm) -> None: @@ -97,17 +96,9 @@ def step(self, length: float_nm) -> None: self.goto(x1, y1) -class BinaryEdgeSweeper(EdgeSweeper): +class BinarySweeper(Sweeper): """A stage-state descriptor used to binary-search the grid edge.""" - 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 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])) @@ -127,31 +118,38 @@ def sweep(self) -> None: def star_sweep( - arms: Literal[3, 4, 5, 6, 7] = 5, - order: Literal[1, 2, 3, 4, 5] = 5, + arms: Literal[3, 4, 5, 6, 7] = 3, + order: Literal[1, 2, 3, 4, 5] = 3, offset: float_deg = 0, ) -> Generator[Vector2, None, None]: - """Sweep window, yielding (x, y) points on its edge as they are found.""" + """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) - _ = EdgeSweeperTeam(name=team) + _ = SweeperTeam(name=team) - # define and sweep with initial sweepers to approximate grid center (order=1) + # define headings, sweep with initial sweepers to approx grid center (order=1) headings = offset + np.linspace(0, 360, num=arms, endpoint=False, dtype=float) - directions = [versor(deg=h) for h in headings] - bes = [BinaryEdgeSweeper(origin=center, heading=d, team=team) for d in directions] - - for sweeper in bes: - sweeper.sweep() - yield sweeper.position + 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 - finished_bes = bes for _ in range(1, order): - new_bes = [] - for a, b in pairwise(finished_bes, closed=True): + new_sweepers = [] + for a, b in pairwise(done_sweepers, closed=True): ns = a.breed(b) ns.sweep() yield ns.position - new_bes.append(ns) - finished_bes = list(chain.from_iterable(zip(finished_bes, new_bes))) + new_sweepers.append(ns) + done_sweepers = list(chain.from_iterable(zip(done_sweepers, new_sweepers))) From 18b5b47362e2fcefe4fe974af0bf49a8e0354409 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Tcho=C5=84?= Date: Wed, 2 Sep 2026 16:15:45 +0200 Subject: [PATCH 7/7] Optimize `grid.py`, `window.py`; remove unused code --- src/instamatic/grid/grid.py | 73 +++-------------------------------- src/instamatic/grid/window.py | 29 ++++---------- 2 files changed, 13 insertions(+), 89 deletions(-) diff --git a/src/instamatic/grid/grid.py b/src/instamatic/grid/grid.py index 70db07aa..6dd85264 100644 --- a/src/instamatic/grid/grid.py +++ b/src/instamatic/grid/grid.py @@ -7,7 +7,7 @@ from typing_extensions import Self from instamatic._collections import NoOverwriteDict -from instamatic._typing import float_nm +from instamatic._typing import float_deg, float_nm from instamatic.grid import versor from instamatic.grid.window import ( GridablePolygonWindow, @@ -18,8 +18,8 @@ from instamatic.utils.pairing import hulam2uv, ij2ulam, ulam2ij, uv2hulam DualIndex = tuple[int, int] -SpiralIndex = Annotated[int, 'positive'] -WindowIndex = Union[DualIndex, SpiralIndex] +SingleIndex = Annotated[int, 'positive'] +WindowIndex = Union[DualIndex, SingleIndex] WindowType = TypeVar('WindowType', bound=GridablePolygonWindow) @@ -31,8 +31,7 @@ class PairingInverse(Protocol): def __call__(self, n: int, /) -> tuple[int, int]: ... -WindowGeometryTuple = tuple[float_nm, float_nm, float, float_nm, Optional[float_nm]] -WindowShapeTuple = tuple[float, float_nm, Optional[float_nm]] +WindowGeometryTuple = tuple[float_nm, float_nm, float_deg, float_nm, Optional[float_nm]] class PeriodicConvexPolygonGrid(Generic[WindowType]): @@ -59,7 +58,7 @@ 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, # signed angle from the X-axis towards a-vector in degrees + 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 @@ -308,65 +307,3 @@ class SquareGrid(PeriodicConvexPolygonGrid[SquareWindow]): GRID_REGISTRY['hexagonal'] = HexagonalGrid GRID_REGISTRY['rectangular'] = RectangularGrid GRID_REGISTRY['square'] = SquareGrid - - -if __name__ == '__main__': - import numpy as np - - residuals = { - 0: np.array( - [ - [572543.92272, 458611.73068], - [564971.6688, 458611.73068], - [564971.6688, 458611.73068], - [564971.6688, 458611.73068], - [552876.06528, 458609.70234], - [552876.06528, 458609.70234], - [538110.39744, 458608.68817], - [538110.39744, 458608.68817], - [538110.39744, 458608.68817], - [518853.20256, 458610.71651], - [518853.20256, 458610.71651], - [496719.8544, 458613.75902], - [496719.8544, 458613.75902], - [487814.08368, 434983.59802], - [484349.97072, 421938.32931], - [484349.97072, 421938.32931], - [483886.27056, 411164.8014], - [483886.27056, 411164.8014], - [486744.23952, 397380.20276], - [486744.23952, 397380.20276], - [486744.23952, 397380.20276], - [486183.55632, 380419.22368], - [486183.55632, 380419.22368], - [496566.80304, 365588.0016], - [496566.80304, 365588.0016], - [503120.73504, 365593.07245], - [503120.73504, 365593.07245], - [511500.67584, 365799.96313], - [511500.67584, 365799.96313], - [520806.5016, 366044.3781], - [520806.5016, 366044.3781], - [532121.69472, 366439.9044], - [543770.26704, 366439.9044], - [543770.26704, 366439.9044], - [555862.83984, 366764.4388], - [555862.83984, 366764.4388], - [568560.04128, 367445.96104], - [568560.04128, 367445.96104], - [572369.65632, 406690.28336], - [576968.77392, 419188.91444], - [575013.95952, 431882.26616], - [572419.6632, 441778.53702], - [572422.69392, 454949.56281], - ] - ) - } - a = SquareGrid(10000, 20000, 0, 50000) - print(a.guess(residuals).to_params()) - print(a.to_params()) - a.refine(residuals) - print(a.to_params()) - a = a.guess(residuals) - a.refine(residuals) - print(a.to_params()) diff --git a/src/instamatic/grid/window.py b/src/instamatic/grid/window.py index f4b0b00c..d1481ada 100644 --- a/src/instamatic/grid/window.py +++ b/src/instamatic/grid/window.py @@ -1,21 +1,15 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Optional, Union +from typing import Optional import numpy as np from typing_extensions import Self -from instamatic._typing import float_nm +from instamatic._typing import float_deg, float_nm from instamatic.grid import versor from instamatic.utils.iterating import pairwise -X = np.array([1, 0], dtype=float) -Y = np.array([0, 1], dtype=float) - - -WindowGeometryTuple = tuple[float_nm, float_nm, float, float_nm, Optional[float_nm]] - class Window(ABC): """Describes an arbitrary single window on a TEM grid.""" @@ -76,7 +70,7 @@ class GridablePolygonWindow(ConvexPolygonWindow): - w: double the distance between window center and its' edge midpoint; """ - INTERIOR_ANGLE: float = ... # class attribute: angle between a and b axes + 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 @@ -85,14 +79,14 @@ def __init__( self, x: float_nm, y: float_nm, - t: float, + 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 = float(t) + 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) @@ -102,14 +96,7 @@ def __init__( def __repr__(self) -> str: """Accurate representation, show params as floats (from to_params).""" - p = self.to_params() - parts = [f'{k}={float(v)}' for k, v in p.items()] - return f'{type(self).__name__}(' + ', '.join(parts) + ')' - - def __str__(self) -> str: - """Nicely display self, show params as integers (from to_params).""" - p = self.to_params() - parts = [f'{k}={int(np.rint(float(v)))}' for k, v in p.items()] + parts = [f'{k}={float(v)}' for k, v in self.to_params().items()] return f'{type(self).__name__}(' + ', '.join(parts) + ')' @property @@ -189,7 +176,7 @@ def to_params(self) -> dict[str, float]: ... class HexagonalWindow(GridablePolygonWindow): """A regular hexagonal window with a 2D "ab" coordinate system.""" - INTERIOR_ANGLE: float = 60.0 + INTERIOR_ANGLE: float_deg = 60.0 ROT60MAT = np.array([[1, -np.sqrt(3)], [np.sqrt(3), 1]], dtype=float) / 2 USES_HEIGHT = False @@ -223,7 +210,7 @@ class RectangularWindow(GridablePolygonWindow): - theta: signed angle from X axis towards A axis and the X-aligned edge. """ - INTERIOR_ANGLE: float = 90.0 + INTERIOR_ANGLE: float_deg = 90.0 USES_HEIGHT = True def __init__(self, *args, **kwargs) -> None: