diff --git a/src/quantem/diffraction/__init__.py b/src/quantem/diffraction/__init__.py index 166c4672b..e26249045 100644 --- a/src/quantem/diffraction/__init__.py +++ b/src/quantem/diffraction/__init__.py @@ -6,6 +6,7 @@ from quantem.diffraction.phase import PhaseMap as PhaseMap from quantem.diffraction.reverse_monte_carlo import ReverseMonteCarlo as ReverseMonteCarlo from quantem.diffraction.strain import StrainMap as StrainMap +from quantem.diffraction.systematic_row import SystematicRowMap as SystematicRowMap from quantem.diffraction import bloch as bloch from quantem.diffraction import calibration as calibration from quantem.diffraction import crystal_rotations as crystal_rotations diff --git a/src/quantem/diffraction/crystal.py b/src/quantem/diffraction/crystal.py index ddd19cf9d..11229e343 100644 --- a/src/quantem/diffraction/crystal.py +++ b/src/quantem/diffraction/crystal.py @@ -450,6 +450,9 @@ def from_cif(cls, file_path: str | Path, name: str | None = None, **kwargs) -> " atoms = read(file_path) assert isinstance(atoms, Atoms) + # the reader adds an ase Spacegroup object; Crystal finds the space + # group itself (Crystal.spacegroup), so we keep only native data + atoms.info.pop("spacegroup", None) return cls(_expand_partial_occupancy(atoms), name=name, **kwargs) @property diff --git a/src/quantem/diffraction/systematic_row.py b/src/quantem/diffraction/systematic_row.py new file mode 100644 index 000000000..aeb3bcc71 --- /dev/null +++ b/src/quantem/diffraction/systematic_row.py @@ -0,0 +1,786 @@ +"""Phase and in-plane orientation from single systematic rows of reflections.""" + +from __future__ import annotations + +from typing import Any + +import matplotlib.pyplot as plt +import numpy as np +from numpy.typing import NDArray + +from quantem.core.io.serialize import AutoSerialize +from quantem.core.visualization.visualization_utils import add_scalebar_to_ax, array_to_rgba +from quantem.diffraction.crystal import Crystal +from quantem.diffraction.orientation import scan_scalebar +from quantem.diffraction.orientation_visualization import phase_color_cycle + + +class SystematicRowMap(AutoSerialize): + """Phase and in-plane orientation mapping from systematic rows. + + A pattern containing a single systematic row of reflections, such as the + local Fourier transform of an HRTEM image showing one set of lattice + fringes, determines the length and the in-plane direction of the + reciprocal lattice vector g, but not the rotation of the crystal about g. + The orientation is then known only up to a fiber about g, and matching + against full zone axis patterns has no unique answer. This class indexes + each row on its own: Friedel pairs of peaks (+g and -g) are combined into + one row, the length of g is assigned to the nearest ring of the candidate + crystals, and the phase, the reflection family and the in-plane angle of + g are stored for up to `num_rows` rows at every position. A position can + hold several rows, from overlapping grains or from a grain showing two + fringe sets. + + The in-plane angle of g is ``atan2(g_col, g_row)``, measured from the scan + row axis toward the column axis, the azimuth convention of the quantEM + calibration and orientation matching, and wrapped to [0, 180) degrees + since g and -g give the same fringes. The lattice fringes run + perpendicular to g. + + Use :meth:`from_vectors` to construct an instance, then :meth:`index`. + + Parameters + ---------- + peaks : Vector + Calibrated peaks with fields (qx, qy, intensity) in 1/Angstroms, + origin at zero. + crystals : list of Crystal + Candidate phases with structure factors calculated. + """ + + _token = object() + + def __init__( + self, + peaks, + crystals: list[Crystal], + k_min: float, + k_max: float, + min_ring_intensity: float, + _token: object | None = None, + ): + if _token is not self._token: + raise RuntimeError("Use SystematicRowMap.from_vectors() to instantiate this class.") + super().__init__() + self.peaks = peaks + self.crystals = list(crystals) + self.k_min = float(k_min) + self.k_max = float(k_max) + self.min_ring_intensity = float(min_ring_intensity) + + # ring table, one entry per reflection family of every crystal + self.ring_k, self.ring_phase, self.ring_label = _ring_table( + self.crystals, self.k_min, self.k_max, self.min_ring_intensity + ) + + # per-row results, (scan_row, scan_col, num_rows), filled by index() + self.k: NDArray | None = None + self.angle_deg: NDArray | None = None + self.intensity: NDArray | None = None + self.phase: NDArray | None = None + self.ring: NDArray | None = None + self.ambiguous: NDArray | None = None + self.metadata: dict[str, Any] = {} + + @classmethod + def from_vectors( + cls, + peaks, + crystals: Crystal | list[Crystal], + k_min: float = 0.1, + k_max: float | None = None, + min_ring_intensity: float = 0.01, + ) -> "SystematicRowMap": + """Bind calibrated peaks to the candidate crystals. + + Parameters + ---------- + peaks : Vector + Calibrated peaks (qx, qy, intensity) in 1/Angstroms, for example + from :meth:`BraggVectors.apply_calibration`. + crystals : Crystal or list of Crystal + Candidate phases, structure factors calculated. + k_min : float, default=0.1 + Peaks closer than this to the origin (1/Angstroms) are ignored, + which removes the direct beam and low-frequency contrast. + k_max : float, optional + Largest ring radius considered (1/Angstroms). None (default) uses + the largest measured peak radius. + min_ring_intensity : float, default=0.01 + A reflection family is a candidate when its structure factor + intensity reaches this fraction of the strongest reflection of its + crystal. + + Returns + ------- + SystematicRowMap + """ + if isinstance(crystals, Crystal): + crystals = [crystals] + if k_max is None: + flat = peaks.select_fields("qx", "qy").numpy() + k_max = float(np.hypot(flat[:, 0], flat[:, 1]).max()) if flat.size else 1.0 + return cls( + peaks=peaks, + crystals=crystals, + k_min=k_min, + k_max=k_max, + min_ring_intensity=min_ring_intensity, + _token=cls._token, + ) + + @property + def names(self) -> list[str]: + """Names of the candidate crystals.""" + return [str(getattr(c, "name", f"phase {i}")) for i, c in enumerate(self.crystals)] + + @property + def shape(self) -> tuple[int, int]: + """Scan shape.""" + return (int(self.peaks.shape[0]), int(self.peaks.shape[1])) + + def rings(self, phase=None) -> list[tuple[str, str, float]]: + """Candidate rings as (phase name, reflection family, radius in 1/Angstroms).""" + idx = range(len(self.ring_k)) if phase is None else self._ring_indices(phase) + return [ + (self.names[self.ring_phase[i]], self.ring_label[i], float(self.ring_k[i])) + for i in idx + ] + + def index( + self, + k_tol: float = 0.01, + pair_tol: float | None = None, + min_intensity: float = 0.0, + num_rows: int = 4, + require_pair: bool = True, + ) -> "SystematicRowMap": + """Pair the peaks into rows and index every row against the rings. + + Each peak is paired with the brightest unpaired peak within + `pair_tol` of its inversion through the origin, and the pair becomes + one row with ``g = (q_1 - q_2) / 2`` and the summed intensity. The row + is assigned to the nearest candidate ring within `k_tol` of ``|g|``; a + row with a second ring inside `k_tol` is marked ambiguous, and a row + with no ring inside `k_tol` keeps phase -1. Rows are sorted by + intensity and the strongest `num_rows` are kept. + + Parameters + ---------- + k_tol : float, default=0.01 + Radial tolerance for the ring assignment, 1/Angstroms. Rings of + different phases closer than about twice this value cannot be + told apart. + pair_tol : float, optional + Largest distance between a peak and the inverted position of its + Friedel mate, 1/Angstroms. None (default) uses `k_tol`. + min_intensity : float, default=0.0 + Rows with a smaller summed intensity are dropped. + num_rows : int, default=4 + Maximum number of rows kept per position. + require_pair : bool, default=True + Keep only rows seen at both +g and -g. The local Fourier transform + of a real image always has both; for 4DSTEM data, where one side + of a row can be weak, set False to keep single peaks as rows. + + Returns + ------- + SystematicRowMap + Self, with the per-row arrays :attr:`k`, :attr:`angle_deg`, + :attr:`intensity`, :attr:`phase`, :attr:`ring` and + :attr:`ambiguous` of shape ``(scan_row, scan_col, num_rows)``. + Empty slots have phase -1, ring -1 and intensity 0. + """ + pair_tol = float(k_tol if pair_tol is None else pair_tol) + R, C = self.shape + M = int(num_rows) + k = np.full((R, C, M), np.nan) + ang = np.full((R, C, M), np.nan) + inten = np.zeros((R, C, M)) + phase = np.full((R, C, M), -1, dtype=np.int16) + ring = np.full((R, C, M), -1, dtype=np.int16) + ambiguous = np.zeros((R, C, M), dtype=bool) + + flat = self.peaks.select_fields("qx", "qy", "intensity").numpy().astype(np.float64) + counts = np.asarray(self.peaks.row_counts(), dtype=int) + cells = np.split(flat, np.cumsum(counts)[:-1]) + for n, data in enumerate(cells): + if data.shape[0] == 0: + continue + rows = _pair_rows(data, self.k_min, pair_tol, require_pair) + if rows.shape[0] == 0: + continue + rows = rows[rows[:, 2] >= min_intensity] + rows = rows[np.argsort(-rows[:, 2])][:M] + r, c = divmod(n, C) + m = rows.shape[0] + kk = np.hypot(rows[:, 0], rows[:, 1]) + k[r, c, :m] = kk + ang[r, c, :m] = np.rad2deg(np.arctan2(rows[:, 1], rows[:, 0])) % 180.0 + inten[r, c, :m] = rows[:, 2] + if self.ring_k.size: + dist = np.abs(kk[:, None] - self.ring_k[None, :]) + order = np.argsort(dist, axis=1) + best = order[:, 0] + hit = dist[np.arange(m), best] <= k_tol + ring[r, c, :m] = np.where(hit, best, -1) + phase[r, c, :m] = np.where(hit, self.ring_phase[best], -1) + if self.ring_k.size > 1: + second = dist[np.arange(m), order[:, 1]] + ambiguous[r, c, :m] = hit & (second <= k_tol) + + self.k, self.angle_deg, self.intensity = k, ang, inten + self.phase, self.ring, self.ambiguous = phase, ring, ambiguous + self.metadata["index"] = dict( + k_tol=float(k_tol), + pair_tol=pair_tol, + min_intensity=float(min_intensity), + num_rows=M, + require_pair=bool(require_pair), + ) + return self + + def _require_index(self) -> None: + if self.phase is None: + raise RuntimeError("Run index() first.") + + def _phase_index(self, phase) -> int: + if isinstance(phase, str): + return self.names.index(phase) + return int(phase) + + def _ring_indices(self, phase) -> list[int]: + p = self._phase_index(phase) + return [i for i in range(len(self.ring_k)) if self.ring_phase[i] == p] + + def _ring_index(self, phase, ring) -> int: + """Ring table index from a family label such as "(200)" or an index.""" + if isinstance(ring, str): + label = ring if ring.startswith("(") else f"({ring})" + for i in self._ring_indices(phase): + if self.ring_label[i] == label: + return i + raise ValueError(f"No ring {ring} for {self.names[self._phase_index(phase)]}.") + return int(ring) + + def _panels(self, phase, ring) -> list[tuple[int, int | None, str]]: + """(phase index, ring index or None, title) of every panel of a plot.""" + if phase is None: + phases = list(range(len(self.crystals))) + elif isinstance(phase, (str, int)): + phases = [self._phase_index(phase)] + else: + phases = [self._phase_index(p) for p in phase] + panels = [] + for p in phases: + name = self.names[p] + if ring is None: + panels.append((p, None, f"{name}, all reflections")) + continue + if isinstance(ring, str) and ring in ("strongest", "all"): + sig = { + g: float((self.intensity * self._row_mask(p, g)).sum()) + for g in self._ring_indices(p) + } + found = [g for g in sig if sig[g] > 0] + if not found: + panels.append((p, None, f"{name}, not found")) + continue + rings = [max(found, key=sig.get)] if ring == "strongest" else found + else: + rings = [self._ring_index(p, ring)] + panels += [(p, g, f"{name} {self.ring_label[g]}") for g in rings] + return panels + + def _row_mask(self, phase=None, ring=None, include_ambiguous: bool = False) -> NDArray: + """Rows of a phase (and ring), as a (R, C, M) boolean mask.""" + self._require_index() + if phase is None: + m = self.phase >= 0 + else: + m = self.phase == self._phase_index(phase) + if ring is not None: + m &= self.ring == self._ring_index(phase, ring) + if not include_ambiguous: + m &= ~self.ambiguous + return m + + def phase_signal(self, include_ambiguous: bool = False) -> NDArray: + """``(scan_row, scan_col, num_phases)`` summed row intensity of every phase. + + Parameters + ---------- + include_ambiguous : bool, default=False + Count rows whose length matches rings of more than one phase. + """ + self._require_index() + out = np.zeros(self.shape + (len(self.crystals),)) + for p in range(len(self.crystals)): + m = self._row_mask(p, include_ambiguous=include_ambiguous) + out[..., p] = (self.intensity * m).sum(axis=-1) + return out + + def phase_index(self, include_ambiguous: bool = False) -> NDArray: + """``(scan_row, scan_col)`` phase with the largest row intensity, -1 for none.""" + sig = self.phase_signal(include_ambiguous) + out = np.argmax(sig, axis=-1).astype(int) + out[sig.max(axis=-1) <= 0] = -1 + return out + + def phase_fractions(self) -> dict[str, float]: + """Fraction of the indexed positions assigned to each phase.""" + idx = self.phase_index() + n = max(int((idx >= 0).sum()), 1) + return {name: float((idx == p).sum()) / n for p, name in enumerate(self.names)} + + def in_plane_angle(self, phase, ring=None) -> tuple[NDArray, NDArray]: + """In-plane angle of g of the strongest row of a phase at every position. + + Parameters + ---------- + phase : str or int + Crystal name or index. + ring : str or int, optional + Restrict to one reflection family, e.g. "200". None (default) + takes the strongest row of any family. + + Returns + ------- + angle_deg : np.ndarray + ``(scan_row, scan_col)`` angle in [0, 180) degrees, NaN where the + phase has no row. + intensity : np.ndarray + ``(scan_row, scan_col)`` intensity of that row, 0 where none. + """ + m = self._row_mask(phase, ring) + w = np.where(m, self.intensity, -np.inf) + j = np.argmax(w, axis=-1) + has = np.isfinite(np.take_along_axis(w, j[..., None], axis=-1)[..., 0]) + angle = np.take_along_axis(self.angle_deg, j[..., None], axis=-1)[..., 0] + inten = np.take_along_axis(self.intensity, j[..., None], axis=-1)[..., 0] + return np.where(has, angle, np.nan), np.where(has, inten, 0.0) + + def num_directions(self, min_angle_deg: float = 15.0) -> NDArray: + """``(scan_row, scan_col)`` number of distinct row directions at each position. + + Rows within `min_angle_deg` of a stronger row count as the same + direction. Positions with two or more directions carry enough + information to determine a full orientation. + """ + self._require_index() + out = np.zeros(self.shape, dtype=int) + for r, c in zip(*np.nonzero(self.intensity[..., 0] > 0)): + dirs: list[float] = [] + for a in self.angle_deg[r, c][self.intensity[r, c] > 0]: + if all(_angle_diff(a, b) > min_angle_deg for b in dirs): + dirs.append(a) + out[r, c] = len(dirs) + return out + + @property + def scan_scalebar(self) -> dict | None: + """Real-space scale bar of the scan, carried by the peaks.""" + return scan_scalebar(getattr(self.peaks, "metadata", {}) or {}) + + def _shade(self, values: NDArray, shade_range, gamma: float) -> NDArray: + """Brightness in [0, 1] from row intensities.""" + v = np.asarray(values, dtype=float) + if isinstance(shade_range, str): + pos = v[v > 0] + # the median signal is full brightness, so half the positions + # saturate and weak fringes fade + hi = float(np.median(pos)) if pos.size else 1.0 + shade_range = (0.0, hi) + lo, hi = shade_range + return np.clip((v - lo) / max(hi - lo, 1e-12), 0.0, 1.0) ** gamma + + def _finish_map(self, ax, rgb, title, scalebar) -> None: + ax.imshow(rgb, interpolation="nearest") + ax.set_title(title) + ax.set_axis_off() + scalebar = self.scan_scalebar if scalebar == "auto" else scalebar + if scalebar: + add_scalebar_to_ax( + ax, + array_size=rgb.shape[1], + sampling=scalebar.get("sampling", 1.0), + length_units=scalebar.get("length", None), + units=scalebar.get("units", "pixels"), + width_px=rgb.shape[0] / 40, + pad_px=rgb.shape[0] / 80, + color=scalebar.get("color", "white"), + loc="lower right", + ) + + def plot_phase( + self, + shade_range="auto", + shade_gamma: float = 1.0, + phase_colors=None, + scalebar="auto", + figax=None, + figsize: tuple[float, float] = (7, 6), + ): + """Map of the phase with the largest row intensity at each position. + + Color gives the phase, brightness the summed row intensity of that + phase. Positions without an indexed row are black. + + Parameters + ---------- + shade_range : tuple or "auto", default="auto" + Intensities mapped to black ... full color; "auto" uses 0 to the + median over the positions with indexed fringes. + shade_gamma : float, default=1.0 + Exponent applied to the brightness. + phase_colors : sequence, optional + One color per crystal; defaults to the quantEM phase palette. + scalebar : dict, "auto" or None, default="auto" + Real-space scale bar; "auto" takes the scan step carried by the + peaks. + figax : tuple, optional + Existing (fig, ax). + figsize : tuple, default=(7, 6) + Figure size in inches. + + Returns + ------- + fig, ax + """ + colors = phase_color_cycle(len(self.crystals), phase_colors) + sig = self.phase_signal() + idx = self.phase_index() + shade = self._shade(sig.max(axis=-1), shade_range, shade_gamma) + rgb = np.where(idx[..., None] >= 0, colors[np.maximum(idx, 0)], 0.0) * shade[..., None] + + fig, ax = plt.subplots(figsize=figsize) if figax is None else figax + self._finish_map(ax, rgb, "phase", scalebar) + frac = self.phase_fractions() + handles = [ + plt.Rectangle((0, 0), 1, 1, color=colors[p], label=f"{n} ({100 * frac[n]:.0f}%)") + for p, n in enumerate(self.names) + ] + ax.legend( + handles=handles, + loc="upper center", + bbox_to_anchor=(0.5, 0.0), + ncol=len(handles), + frameon=False, + ) + return fig, ax + + def plot_orientation( + self, + phase=None, + ring="strongest", + angle_range: tuple[float, float] | None = None, + shade_range="auto", + shade_gamma: float = 1.0, + scalebar="auto", + legend: bool = True, + axsize: tuple[float, float] = (6, 6), + ): + """Maps of the in-plane angle of g, one per phase and reflection family. + + Hue gives the angle of the strongest row of the phase and family, + brightness its intensity. Each title names the phase and the + reflection family. By default the hue runs once around the color wheel over + [0, 180) degrees, so every angle has its own color; for a textured + sample, `angle_range` spreads the colors over a narrower window. + + Parameters + ---------- + phase : str, int or list, optional + Phase(s) to show. None (default) shows every phase. + ring : str, int or None, default="strongest" + Reflection family of each map. "strongest" takes the family with + the largest summed row intensity of each phase, "all" draws one + map per family with indexed rows, a label such as "200" or a ring + index selects that family, and None takes the strongest row of + any family. + angle_range : tuple, optional + ``(low, high)`` angles in degrees mapped to the ends of the + "turbo" colormap. Angles are wrapped modulo 180 into the half + turn centered on the window, and values outside it are clipped. + None (default) uses the cyclic color wheel over [0, 180). + shade_range : tuple or "auto", default="auto" + Intensities mapped to black ... full color, per map; "auto" uses 0 + to the median over the positions with indexed fringes. + shade_gamma : float, default=1.0 + Exponent applied to the brightness. + scalebar : dict, "auto" or None, default="auto" + Real-space scale bar. + legend : bool, default=True + Draw the color wheel of the angle. + axsize : tuple, default=(6, 6) + Size of each map in inches. + + Returns + ------- + fig, axs + """ + panels = self._panels(phase, ring) + fig, axs = plt.subplots( + 1, len(panels), figsize=(axsize[0] * len(panels), axsize[1]), squeeze=False + ) + for ax, (p, g, title) in zip(axs[0], panels): + angle, inten = self.in_plane_angle(p, g) + rgb = _angle_rgb(angle, angle_range, self._shade(inten, shade_range, shade_gamma)) + self._finish_map(ax, rgb, title, scalebar) + if legend: + fig.subplots_adjust(right=0.88) + if angle_range is None: + _angle_legend(fig.add_axes((0.89, 0.7, 0.1, 0.2))) + else: + cax = fig.add_axes((0.9, 0.25, 0.015, 0.5)) + lo, hi = angle_range + t = np.linspace(lo, hi, 256) + cax.imshow( + _angle_rgb(t, angle_range)[::-1, None], aspect="auto", extent=(0, 1, lo, hi) + ) + cax.set_xticks([]) + cax.yaxis.tick_right() + cax.set_title("g angle", fontsize=9) + return fig, axs + + def plot_angle_distribution( + self, + phase=None, + ring="strongest", + angle_range: tuple[float, float] | None = None, + bin_deg: float = 1.0, + weight: str = "intensity", + axsize: tuple[float, float] = (5, 3), + ): + """Histograms of the in-plane angle of g, one per phase and reflection family. + + The bars take the colors of :meth:`plot_orientation` with the same + `angle_range`, so each histogram also serves as the color legend of + the matching map. A narrow distribution is a fiber texture: the + crystals share the direction of g in the image plane, and their + rotation about g is not measured. + + Parameters + ---------- + phase : str, int or list, optional + Phase(s) to show. None (default) shows every phase. + ring : str, int or None, default="strongest" + Reflection family of each histogram, as in + :meth:`plot_orientation`; None pools every family of a phase. + angle_range : tuple, optional + ``(low, high)`` window of the angle axis in degrees, with angles + wrapped modulo 180 into the half turn centered on it. None + (default) shows [0, 180). + bin_deg : float, default=1.0 + Histogram bin width in degrees. + weight : {"intensity", "count"}, default="intensity" + Weight each row by its intensity, or count rows alike. + axsize : tuple, default=(5, 3) + Size of each panel in inches. + + Returns + ------- + fig, axs + """ + panels = self._panels(phase, ring) + fig, axs = plt.subplots( + 1, len(panels), figsize=(axsize[0] * len(panels), axsize[1]), squeeze=False + ) + lo, hi = (0.0, 180.0) if angle_range is None else (float(v) for v in angle_range) + mid = 0.5 * (lo + hi) + edges = np.arange(lo, hi + 0.5 * bin_deg, bin_deg) + centers = 0.5 * (edges[:-1] + edges[1:]) + for ax, (p, g, title) in zip(axs[0], panels): + m = self._row_mask(p, g) + a = self.angle_deg[m] + if angle_range is not None: + a = (a - mid + 90.0) % 180.0 - 90.0 + mid + w = self.intensity[m] if weight == "intensity" else np.ones(a.size) + h, _ = np.histogram(a, edges, weights=w) + h = h / max(w.sum(), 1e-12) + ax.bar(centers, h, width=bin_deg, color=_angle_rgb(centers, angle_range), lw=0) + ax.set_xlim(lo, hi) + ax.set_xlabel("g angle (degrees)") + ax.set_ylabel("fraction of count" if weight == "count" else "fraction of intensity") + ax.set_title(title) + fig.tight_layout() + return fig, axs + + def plot_rings( + self, + bin_size: float = 0.002, + figsize: tuple[float, float] = (10, 4), + ): + """Histogram of the row lengths under the candidate rings. + + Checks the calibration and the ring assignment: every peak of the + histogram should sit on a ring, and rings of different phases closer + than twice `k_tol` are ambiguous. + + Parameters + ---------- + bin_size : float, default=0.002 + Histogram bin width, 1/Angstroms. + figsize : tuple, default=(10, 4) + Figure size in inches. + + Returns + ------- + fig, ax + """ + self._require_index() + colors = phase_color_cycle(len(self.crystals)) + valid = self.intensity > 0 + k = self.k[valid] + edges = np.arange(self.k_min, self.k_max + bin_size, bin_size) + fig, ax = plt.subplots(figsize=figsize) + h_all, _ = np.histogram(k, edges) + ax.stairs(h_all, edges, color="0.6", fill=True, label="unassigned") + bottom = np.zeros_like(h_all, dtype=float) + for p, name in enumerate(self.names): + h, _ = np.histogram(self.k[valid & (self.phase == p) & ~self.ambiguous], edges) + ax.stairs(bottom + h, edges, baseline=bottom, color=colors[p], fill=True, label=name) + bottom = bottom + h + if self.ambiguous.any(): + h, _ = np.histogram(self.k[valid & self.ambiguous], edges) + ax.stairs(bottom + h, edges, baseline=bottom, color="k", fill=True, label="ambiguous") + k_tol = self.metadata["index"]["k_tol"] + top = max(h_all.max(), 1) + for i, kr in enumerate(self.ring_k): + c = colors[self.ring_phase[i]] + ax.axvspan(kr - k_tol, kr + k_tol, color=c, alpha=0.15, lw=0) + ax.axvline(kr, color=c, lw=1) + ax.text( + kr, + 1.12 * top, + f"{self.names[self.ring_phase[i]]} {self.ring_label[i]}", + color="k", + ha="center", + va="bottom", + rotation=90, + fontsize=8, + clip_on=False, + ) + ax.set_xlim(self.k_min, self.k_max) + ax.set_ylim(0, 1.1 * top) + ax.set_xlabel("|g| (1/Å)") + ax.set_ylabel("count") + ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), frameon=False) + fig.tight_layout() + return fig, ax + + def __repr__(self) -> str: + head = f"SystematicRowMap({self.shape[0]} x {self.shape[1]}, phases {self.names}" + if self.phase is None: + return head + ")" + n = int((self.intensity > 0).sum()) + assigned = int((self.phase >= 0).sum()) + amb = int(self.ambiguous.sum()) + return head + f", {n} rows, {assigned} indexed, {amb} ambiguous)" + + +def _ring_table(crystals, k_min: float, k_max: float, min_ring_intensity: float): + """Ring radius, phase index and family label of every candidate ring.""" + ks, phases, labels = [], [], [] + for p, xtl in enumerate(crystals): + if getattr(xtl, "g_len", None) is None: + raise ValueError("Calculate the structure factors of every crystal first.") + g = xtl.g_len.numpy() + f = xtl.struct_factors_int.numpy() + hkl = xtl.hkl.numpy() + nz = g > 1e-6 + keep = nz & (g >= k_min) & (g <= k_max) & (f >= min_ring_intensity * f[nz].max()) + radii = np.round(g[keep], 4) + for kr in np.unique(radii): + members = np.nonzero(keep)[0][radii == kr] + # label each family by its largest member with non-negative + # indices, e.g. (200) rather than (0-20) + best = max(members, key=lambda i: (bool((hkl[i] >= 0).all()), tuple(hkl[i]))) + ks.append(float(kr)) + phases.append(p) + labels.append("(" + "".join(str(int(v)) for v in hkl[best]) + ")") + order = np.argsort(ks) + return ( + np.asarray(ks)[order], + np.asarray(phases, dtype=int)[order], + [labels[i] for i in order], + ) + + +def _pair_rows(data: NDArray, k_min: float, pair_tol: float, require_pair: bool) -> NDArray: + """Friedel-paired rows (g_row, g_col, intensity) of one pattern.""" + q = data[:, :2] + inten = data[:, 2] + keep = np.hypot(q[:, 0], q[:, 1]) >= k_min + q, inten = q[keep], inten[keep] + order = np.argsort(-inten) + used = np.zeros(q.shape[0], dtype=bool) + rows = [] + for i in order: + if used[i]: + continue + used[i] = True + d = np.hypot(*(q + q[i]).T) + d[used] = np.inf + j = int(np.argmin(d)) if d.size else -1 + if j >= 0 and d[j] <= pair_tol: + used[j] = True + g = 0.5 * (q[i] - q[j]) + rows.append((g[0], g[1], inten[i] + inten[j])) + elif not require_pair: + rows.append((q[i, 0], q[i, 1], inten[i])) + return np.asarray(rows, dtype=float).reshape(-1, 3) + + +def _angle_diff(a: float, b: float) -> float: + """Distance between two angles modulo 180 degrees.""" + d = abs(a - b) % 180.0 + return min(d, 180.0 - d) + + +def _angle_rgb( + angle_deg: NDArray, + angle_range: tuple[float, float] | None = None, + brightness: NDArray | None = None, +) -> NDArray: + """Perceptually uniform colors of in-plane angles; NaN is black. + + The colors come from :func:`array_to_rgba`, which works in the CIECAM02 + JCh space: the angle sets the hue at constant lightness, and `brightness` + sets the lightness, so every angle is equally bright and shading darkens + every hue alike. Without a range the hue runs once around the color + circle over [0, 180) degrees. With a range, angles are wrapped into the + half turn centered on it and mapped linearly onto 300 degrees of the hue + circle, so the two ends keep different colors, clipped beyond the ends. + """ + angle_deg = np.asarray(angle_deg, dtype=float) + finite = np.isfinite(angle_deg) + a = np.nan_to_num(angle_deg) + if angle_range is None: + phase = np.deg2rad(2.0 * a) + else: + lo, hi = (float(v) for v in angle_range) + mid = 0.5 * (lo + hi) + a = (a - mid + 90.0) % 180.0 - 90.0 + mid + t = np.clip((a - lo) / (hi - lo), 0.0, 1.0) + phase = np.deg2rad(300.0 * t - 150.0) + amp = np.ones_like(a) if brightness is None else np.asarray(brightness, dtype=float) + amp = np.where(finite, np.clip(amp, 0.0, 1.0), 0.0) + # array_to_rgba stacks 2D images; run any shape through it as a column + rgba = array_to_rgba(amp.reshape(-1, 1), phase.reshape(-1, 1)) + return rgba[..., :3].reshape(a.shape + (3,)) + + +def _angle_legend(ax) -> None: + """Half-wheel legend of the in-plane angle hue, in the image frame.""" + n = 181 + t = np.linspace(0, 2 * np.pi, 2 * n) + r = np.linspace(0.4, 1.0, 24) + T, Rr = np.meshgrid(t, r) + # image frame: angle from the row axis (down) toward the column axis (right) + x, y = Rr * np.sin(T), -Rr * np.cos(T) + color = _angle_rgb(np.rad2deg(T)) + ax.scatter(x.ravel(), y.ravel(), c=color.reshape(-1, 3), s=2, marker="s", lw=0) + ax.set_aspect("equal") + ax.set_axis_off() + ax.set_title("g angle", fontsize=9) diff --git a/src/quantem/imaging/__init__.py b/src/quantem/imaging/__init__.py index 3637acdc2..0550fb7df 100644 --- a/src/quantem/imaging/__init__.py +++ b/src/quantem/imaging/__init__.py @@ -1,3 +1,4 @@ from quantem.imaging.drift import DriftCorrection as DriftCorrection from quantem.imaging.lattice import Lattice as Lattice from quantem.imaging.lattice_visualization import PLOT_REGISTRY as PLOT_REGISTRY +from quantem.imaging.local_fft import LocalFFT as LocalFFT diff --git a/src/quantem/imaging/local_fft.py b/src/quantem/imaging/local_fft.py new file mode 100644 index 000000000..b159c92a2 --- /dev/null +++ b/src/quantem/imaging/local_fft.py @@ -0,0 +1,502 @@ +from __future__ import annotations + +from typing import Any, Literal + +import matplotlib.pyplot as plt +import numpy as np +import torch +from numpy.typing import NDArray + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.io.serialize import AutoSerialize +from quantem.core.visualization import show_2d + +# image units converted to Angstroms, so the reciprocal axes match Crystal +_LENGTH_TO_ANGSTROM = {"a": 1.0, "å": 1.0, "angstrom": 1.0, "nm": 10.0, "pm": 0.01, "um": 1e4} + + +def window_mask( + window_size: int, + edge_width: float | None = None, +) -> NDArray: + """Smooth circular window for local Fourier transforms. + + The window is 1 inside a disk and falls to 0 at radius ``window_size / 2`` + as ``sin(u * pi / 2) ** 2``, where ``u`` runs from 1 at the inner edge of + the taper to 0 at the outer radius. With the default ``edge_width`` the + taper spans the whole radius, so ``u = 1 - r / R`` and no part of the + window is flat. + + Parameters + ---------- + window_size : int + Width of the square window in pixels; the disk diameter. + edge_width : float, optional + Radial width of the taper in pixels. None (default) uses the full + radius ``window_size / 2``. + + Returns + ------- + np.ndarray + ``(window_size, window_size)`` float32 window. + """ + radius = window_size / 2.0 + edge = radius if edge_width is None else float(edge_width) + if not 0.0 < edge <= radius: + raise ValueError(f"edge_width must be in (0, {radius}], got {edge_width}.") + x = np.arange(window_size) - (window_size - 1) / 2.0 + r = np.hypot(x[:, None], x[None, :]) + u = np.clip((radius - r) / edge, 0.0, 1.0) + return (np.sin(u * np.pi / 2.0) ** 2).astype(np.float32) + + +class LocalFFT(AutoSerialize): + """Pseudo-4DSTEM dataset from the local Fourier transforms of an image. + + A circular window with a smooth edge (:func:`window_mask`) is stepped over + the image on a regular grid, and the Fourier transform of each windowed + region becomes the diffraction pattern at that grid position. Lattice + fringes in the window produce pairs of spots at plus and minus their + reciprocal lattice vectors, so the standard 4DSTEM workflow (Bragg spot + detection, calibration, orientation and phase matching) applies to the + result. Compared with a 4DSTEM experiment, the pattern center is known + exactly, the patterns have no descan or elliptic distortion, and the + reciprocal pixel size follows from the image calibration and the window + size. The intensities are set by the contrast transfer of the microscope + and the spot shape by the window, so they are not kinematic diffraction + intensities. + + Each window has its weighted mean subtracted before the transform, which + removes most of the zero-frequency peak. The image is divided by its + global mean, so the patterns are in units of fractional contrast, and the + power is divided by the sum of the squared window: white noise of variance + ``sigma ** 2`` (in fractional contrast) then gives a flat background of + ``sigma ** 2``. The background of an HRTEM image falls steeply with + spatial frequency (amorphous contrast, the envelope of the contrast + transfer and the detector response), so by default every pattern is + normalized by the azimuthal median of the scan-mean pattern + (:attr:`background`) as ``I / background - 1``. The background is then 0 + and the noise close to 1 at every frequency, the spot intensities are + excess over the local background, and one detection threshold applies + across the whole pattern. + + Use :meth:`from_image` to construct an instance. The pseudo-4DSTEM data is + held in :attr:`dataset`, and :attr:`origins` holds the exact pattern + centers in the form :meth:`BraggVectors.correct_peak_origins` expects. + + Parameters + ---------- + image : Dataset2d + The image, with its sampling converted to Angstroms when the units + are known. + window_size, step, edge_width, k_max, fft_size, intensity, normalize, center_radius + See :meth:`from_image`. + """ + + _token = object() + + def __init__( + self, + image: Dataset2d, + window_size: int, + step: int, + edge_width: float | None, + k_max: float | None, + fft_size: int, + intensity: str, + normalize: bool, + center_radius: float, + _token: object | None = None, + ): + if _token is not self._token: + raise RuntimeError("Use LocalFFT.from_image() to instantiate this class.") + super().__init__() + self.image = image + self.window_size = int(window_size) + self.step = int(step) + self.edge_width = edge_width + self.k_max = k_max + self.fft_size = int(fft_size) + self.intensity = intensity + self.normalize = bool(normalize) + self.center_radius = float(center_radius) + self.window = window_mask(self.window_size, edge_width) + self.dataset: Dataset4dstem | None = None + self.background: NDArray | None = None + self.corners: NDArray | None = None + self.metadata: dict[str, Any] = {} + + @classmethod + def from_image( + cls, + image: Dataset2d | NDArray, + window_size: int = 128, + step: int | None = None, + edge_width: float | None = None, + k_max: float | None = None, + fft_size: int | None = None, + intensity: Literal["power", "amplitude"] = "power", + normalize: bool = True, + center_radius: float = 1.5, + sampling: float | None = None, + device: str = "cpu", + batch_size: int = 1024, + ) -> "LocalFFT": + """Compute the local Fourier transforms of an image. + + The window grid is centered in the image, with ``step`` pixels between + neighboring windows and every window fully inside the image. + + Parameters + ---------- + image : Dataset2d | np.ndarray + The image. A Dataset2d in nm, pm, um or Angstroms is converted to + Angstroms; an array needs ``sampling``. + window_size : int, default=128 + Width of the square window in pixels, and the diameter of the + circular mask. The reciprocal pixel size is + ``1 / (fft_size * sampling)``, and a spot from a crystal filling + the window is about two pixels wide. + step : int, optional + Step between windows in pixels. None (default) uses + ``window_size // 4``. + edge_width : float, optional + Radial width of the ``sin ** 2`` taper in pixels. None (default) + tapers over the full radius; a narrower edge keeps more of the + window at full weight, which sharpens the spots and raises the + side lobes of strong fringes. + k_max : float, optional + Crop each pattern to ``|k_x|, |k_y| <= k_max`` (1/Angstroms) around + the center. None (default) keeps every frequency up to Nyquist. + Spots beyond the information limit only add noise to the + detection and the matching. + fft_size : int, optional + Fourier transform size in pixels; the window is zero padded to it, + which samples the pattern more finely without adding information. + None (default) uses ``window_size``. + intensity : {"power", "amplitude"}, default="power" + Store the power spectrum ``|F| ** 2`` or the amplitude ``|F|``. + The amplitude compresses the range between strong and weak + fringes. + normalize : bool, default=True + Remove :attr:`background`, the azimuthal median of the scan-mean + pattern, as ``I / background - 1``. The background is then 0 and + the noise close to 1 at every frequency, and a spot intensity is + its excess over the local background. The median ignores the + strong arcs of a textured sample. + center_radius : float, default=1.5 + With ``normalize``, pixels closer than this to zero frequency are + set to 0. The orientation matching weights every spot by its + radius, so the center carries no weight there and nothing is + lost by removing it. + sampling : float, optional + Image pixel size in Angstroms. Overrides the calibration of a + Dataset2d, and is required for an array. + device : str, default="cpu" + Torch device for the Fourier transforms. + batch_size : int, default=1024 + Number of windows transformed at once. + + Returns + ------- + LocalFFT + The instance, with :attr:`dataset` holding the pseudo-4DSTEM data + in shape ``(scan_row, scan_col, k_row, k_col)``. + """ + image = _as_angstrom_dataset(image, sampling) + window_size = int(window_size) + step = window_size // 4 if step is None else int(step) + fft_size = window_size if fft_size is None else int(fft_size) + if fft_size < window_size: + raise ValueError("fft_size must be at least window_size.") + if intensity not in ("power", "amplitude"): + raise ValueError("intensity must be 'power' or 'amplitude'.") + if min(image.shape) < window_size: + raise ValueError("The window is larger than the image.") + + lfft = cls( + image=image, + window_size=window_size, + step=step, + edge_width=edge_width, + k_max=k_max, + fft_size=fft_size, + intensity=intensity, + normalize=normalize, + center_radius=center_radius, + _token=cls._token, + ) + lfft._compute(device=device, batch_size=batch_size) + return lfft + + @property + def sampling(self) -> float: + """Image pixel size in Angstroms.""" + return float(self.image.sampling[0]) + + @property + def pixel_size(self) -> float: + """Reciprocal pixel size of the patterns in 1/Angstroms.""" + return 1.0 / (self.fft_size * self.sampling) + + @property + def scan_shape(self) -> tuple[int, int]: + """Number of windows along the image rows and columns.""" + return tuple(int(v) for v in self.dataset.shape[:2]) + + @property + def center(self) -> tuple[float, float]: + """Pattern pixel (row, col) of zero spatial frequency.""" + c = float(self.metadata["center_px"]) + return (c, c) + + @property + def origins(self) -> NDArray: + """``(scan_row, scan_col, 2)`` pattern centers for every window. + + Pass to :meth:`BraggVectors.correct_peak_origins` with + ``origin_ref=(0, 0)`` to put the detected spots in a frame centered on + zero frequency, in place of measuring and fitting the origins. + """ + return np.broadcast_to(np.asarray(self.center), self.scan_shape + (2,)).copy() + + @property + def window_fwhm(self) -> float: + """Diameter of the window where its weight is 1/2, in Angstroms. + + This is the effective spatial resolution of the maps: ``R`` for the + full-radius taper, ``2 R - edge_width`` for a narrower edge, with ``R`` + the window radius. + """ + radius = self.window_size / 2.0 + edge = radius if self.edge_width is None else float(self.edge_width) + return (2.0 * radius - edge) * self.sampling + + @property + def window_centers(self) -> NDArray: + """``(scan_row, scan_col, 2)`` window centers in image pixels.""" + return self.corners + (self.window_size - 1) / 2.0 + + def _compute(self, device: str, batch_size: int) -> None: + """Fill :attr:`dataset` with the windowed Fourier transforms.""" + img = np.asarray(self.image.array, dtype=np.float32) + mean = float(img.mean()) + if mean <= 0: + raise ValueError("The image mean must be positive to normalize the contrast.") + img = img / mean + + n, s, nf = self.window_size, self.step, self.fft_size + num = [(size - n) // s + 1 for size in img.shape] + offset = [(size - n - (m - 1) * s) // 2 for size, m in zip(img.shape, num)] + rows = offset[0] + s * np.arange(num[0]) + cols = offset[1] + s * np.arange(num[1]) + self.corners = np.stack(np.meshgrid(rows, cols, indexing="ij"), axis=-1) + + # crop symmetric about zero frequency to an odd size, so the center is + # a pixel and every spot keeps its Friedel mate + c = nf // 2 + if self.k_max is None: + h = c - 1 if nf % 2 == 0 else c + else: + h = min(int(np.floor(self.k_max / self.pixel_size)), c - 1 if nf % 2 == 0 else c) + crop = slice(c - h, c + h + 1) + size = 2 * h + 1 + + dev = torch.device(device) + w = torch.as_tensor(self.window, device=dev) + w_sum = w.sum() + norm = (w * w).sum() + image_t = torch.as_tensor(img, device=dev) + flat = self.corners.reshape(-1, 2) + out = np.empty((flat.shape[0], size, size), dtype=np.float32) + ar = torch.arange(n, device=dev) + for start in range(0, flat.shape[0], batch_size): + corner = torch.as_tensor(flat[start : start + batch_size], device=dev) + ri = (corner[:, 0, None] + ar[None, :])[:, :, None] + ci = (corner[:, 1, None] + ar[None, :])[:, None, :] + x = image_t[ri, ci] + x_mean = (x * w).sum(dim=(1, 2), keepdim=True) / w_sum + x = (x - x_mean) * w + F = torch.fft.fftshift(torch.fft.fft2(x, s=(nf, nf)), dim=(-2, -1))[:, crop, crop] + p = F.real**2 + F.imag**2 + p = p / norm + if self.intensity == "amplitude": + p = torch.sqrt(p) + out[start : start + batch_size] = p.cpu().numpy() + + # azimuthal median of the mean pattern in 1 pixel shells, interpolated + # linearly at the radius of every pixel so the background is smooth + x = np.arange(size) - h + r = np.hypot(x[:, None], x[None, :]) + r_int = np.round(r).astype(int) + mean_dp = out.mean(axis=0) + profile = np.array([np.median(mean_dp[r_int == i]) for i in range(r_int.max() + 1)]) + self.background = np.interp(r, np.arange(profile.size), profile).astype(np.float32) + if self.normalize: + # inside the center disk the background is the zero-frequency + # remainder of the mean subtraction, close to 0, so the ratio is + # meaningless there; the disk is set to 0 instead + inside = r < self.center_radius + bg = np.where(inside, 1.0, self.background) + out /= bg + out -= 1.0 + out[:, inside] = 0.0 + + dx = self.sampling + dk = self.pixel_size + origin_image = np.asarray(self.image.origin, dtype=float) + self.dataset = Dataset4dstem.from_array( + out.reshape(num[0], num[1], size, size), + name=f"{self.image.name} local FFT", + origin=[ + origin_image[0] + (offset[0] + (n - 1) / 2) * dx, + origin_image[1] + (offset[1] + (n - 1) / 2) * dx, + -h * dk, + -h * dk, + ], + sampling=[s * dx, s * dx, dk, dk], + units=["A", "A", "1/A", "1/A"], + signal_units=( + "ratio to background" + if self.normalize + else ("contrast^2" if self.intensity == "power" else "contrast") + ), + ) + self.metadata.update( + center_px=float(h), + image_mean=mean, + window_size=n, + step=s, + edge_width=self.edge_width, + fft_size=nf, + k_max=self.k_max, + intensity=self.intensity, + normalize=self.normalize, + center_radius=self.center_radius, + ) + + def show( + self, + positions: list[tuple[int, int]] | None = None, + norm: dict | str | None = None, + axsize: tuple[float, float] = (4, 4), + **kwargs, + ): + """Show the window grid on the image, and some windows with their patterns. + + Parameters + ---------- + positions : list of (row, col), optional + Scan positions to show. None (default) shows the center of the + scan. + norm : dict | str, optional + Normalization of the patterns, passed to + :func:`~quantem.core.visualization.show_2d`. None (default) uses + ``{"power": 0.5, "upper_quantile": 0.999}``. + axsize : tuple, default=(4, 4) + Size of each panel in inches. + **kwargs + Further arguments of :func:`~quantem.core.visualization.show_2d` + for the pattern panels. + + Returns + ------- + fig, axs + The matplotlib figure and axes. + """ + if positions is None: + positions = [tuple(v // 2 for v in self.scan_shape)] + norm = {"power": 0.5, "upper_quantile": 0.999} if norm is None else norm + num = len(positions) + fig, axs = plt.subplots(num, 3, figsize=(3 * axsize[0], num * axsize[1]), squeeze=False) + + img = np.asarray(self.image.array, dtype=np.float32) + lo, hi = np.quantile(img, (0.01, 0.99)) + n = self.window_size + for a, (r, c) in enumerate(positions): + ax = axs[a, 0] + ax.imshow(img, cmap="gray", vmin=lo, vmax=hi) + # the area covered by the window centers + (r_lo, c_lo), (r_hi, c_hi) = self.window_centers[0, 0], self.window_centers[-1, -1] + ax.plot( + [c_lo, c_hi, c_hi, c_lo, c_lo], + [r_lo, r_lo, r_hi, r_hi, r_lo], + color=(1.0, 0.3, 0.3), + lw=1, + ) + r0, c0 = self.corners[r, c] + t = np.linspace(0, 2 * np.pi, 181) + ax.plot( + c0 + (n - 1) / 2 + n / 2 * np.cos(t), + r0 + (n - 1) / 2 + n / 2 * np.sin(t), + color=(1.0, 0.8, 0.0), + ) + ax.set_title(f"window ({r}, {c})") + ax.set_axis_off() + + crop = img[r0 : r0 + n, c0 : c0 + n] * self.window + axs[a, 1].imshow(crop, cmap="gray") + axs[a, 1].set_title("masked window") + axs[a, 1].set_axis_off() + + show_2d( + np.asarray(self.dataset.array[r, c]), + norm=norm, + scalebar={"sampling": self.pixel_size, "units": "1/Å"}, + title="local FFT", + figax=(fig, axs[a, 2]), + **kwargs, + ) + fig.tight_layout() + return fig, axs + + def __repr__(self) -> str: + if self.dataset is None: + return "LocalFFT()" + dx = self.sampling + return ( + f"LocalFFT\n" + f" image {self.image.shape[0]} x {self.image.shape[1]} px at {dx:.4f} A\n" + f" window {self.window_size} px = {self.window_size * dx:.1f} A diameter, " + f"{self.window_fwhm:.1f} A FWHM (effective resolution)\n" + f" step {self.step} px = {self.step * dx:.1f} A, " + f"scan {self.scan_shape[0]} x {self.scan_shape[1]}\n" + f" patterns {self.dataset.shape[2]} x {self.dataset.shape[3]} px at " + f"{self.pixel_size:.5f} 1/A per px, |k| <= {self.metadata['center_px'] * self.pixel_size:.3f} 1/A" + ) + + +def _as_angstrom_dataset(image: Dataset2d | NDArray, sampling: float | None) -> Dataset2d: + """A float32 Dataset2d calibrated in Angstroms.""" + if isinstance(image, Dataset2d): + units = [str(u).strip().lower() for u in image.units] + scale = _LENGTH_TO_ANGSTROM.get(units[0]) + array = np.asarray(image.array, dtype=np.float32) + if sampling is not None: + sampling_A = float(sampling) + origin_A = np.zeros(2) + elif scale is not None: + sampling_A = float(image.sampling[0]) * scale + origin_A = np.asarray(image.origin, dtype=float) * scale + else: + raise ValueError( + f"Image units {image.units} are not a length; pass sampling in Angstroms." + ) + if not np.isclose(image.sampling[0], image.sampling[1], rtol=1e-6): + raise ValueError("LocalFFT needs square image pixels.") + name = image.name + else: + if sampling is None: + raise ValueError("Pass sampling (Angstroms per pixel) with an array image.") + array = np.asarray(image, dtype=np.float32) + sampling_A = float(sampling) + origin_A = np.zeros(2) + name = "image" + return Dataset2d.from_array( + array, + name=name, + origin=origin_A, + sampling=[sampling_A, sampling_A], + units=["A", "A"], + ) diff --git a/tests/diffraction/test_crystal.py b/tests/diffraction/test_crystal.py index b8a358330..d33fae279 100644 --- a/tests/diffraction/test_crystal.py +++ b/tests/diffraction/test_crystal.py @@ -392,6 +392,7 @@ def test_from_cif_keeps_partial_occupancy(tmp_path): assert content["Li"] == pytest.approx(8.444, abs=1e-6) assert content["O"] == pytest.approx(12.0, abs=1e-6) assert xtl.spacegroup.startswith("C2/m") + assert "spacegroup" not in xtl.atoms.info # no ase objects inside Crystal def test_pseudo_symmetry_names_the_breaking_reflection(tmp_path): diff --git a/tests/diffraction/test_systematic_row.py b/tests/diffraction/test_systematic_row.py new file mode 100644 index 000000000..49f54d974 --- /dev/null +++ b/tests/diffraction/test_systematic_row.py @@ -0,0 +1,44 @@ +"""SystematicRowMap: phase and in-plane angle from Friedel-paired rows.""" + +import numpy as np +from ase.build import bulk + +from quantem.core.datastructures.vector import Vector +from quantem.diffraction import Crystal, SystematicRowMap + + +def test_row_indexing(): + lif = Crystal.from_ase(bulk("LiF", "rocksalt", a=4.027, cubic=True), name="LiF", verbose=False) + lif.calculate_structure_factors(k_max=1.0) + li = Crystal.from_ase(bulk("Li", "bcc", a=3.51, cubic=True), name="Li", verbose=False) + li.calculate_structure_factors(k_max=1.0) + + def pair(k, deg, inten): + # deg is measured from the row axis toward the column axis + g = k * np.array([np.cos(np.deg2rad(deg)), np.sin(np.deg2rad(deg))]) + return [[g[0], g[1], inten], [-g[0], -g[1], inten]] + + k200 = 2 / 4.027 + k110 = np.sqrt(2) / 3.51 + cells = [ + np.array(pair(k200, 30.0, 5.0)), # LiF (200) at 30 degrees + np.array(pair(k110, 100.0, 4.0) + pair(k200, 10.0, 1.0)), # Li wins, two directions + np.array([[0.3, 0.1, 9.0]]), # unpaired peak, no row + np.array(pair(0.3, 45.0, 2.0)), # paired but on no ring + ] + peaks = Vector.from_data( + [cells[:2], cells[2:]], fields=["qx", "qy", "intensity"], units=["A^-1"] * 2 + ["counts"] + ) + rm = SystematicRowMap.from_vectors(peaks, [lif, li], k_min=0.2, k_max=0.6).index(k_tol=0.01) + + phase = rm.phase_index() + assert phase.tolist() == [[0, 1], [-1, -1]], phase + assert rm.ring_label[rm.ring[0, 0, 0]] == "(200)" + angle, _ = rm.in_plane_angle("LiF") + assert np.isclose(angle[0, 0], 30.0) and np.isclose(angle[0, 1], 10.0), angle + assert rm.num_directions()[0, 1] == 2 + assert rm.intensity[1, 0, 0] == 0, "an unpaired peak must not form a row" + assert rm.phase[1, 1, 0] == -1 and rm.intensity[1, 1, 0] > 0, "off-ring row kept unassigned" + single = SystematicRowMap.from_vectors(peaks, [lif, li], k_min=0.2, k_max=0.6) + single.index(k_tol=0.01, require_pair=False) + assert single.intensity[1, 0, 0] == 9.0, "require_pair=False keeps single peaks" diff --git a/tests/imaging/test_local_fft.py b/tests/imaging/test_local_fft.py new file mode 100644 index 000000000..09748a5cb --- /dev/null +++ b/tests/imaging/test_local_fft.py @@ -0,0 +1,31 @@ +"""LocalFFT: windowed Fourier transforms of an image as a pseudo-4DSTEM dataset.""" + +import numpy as np + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.imaging.local_fft import LocalFFT, window_mask + + +def test_local_fft(): + # full-radius taper sin^2(pi/2 * (1 - r/R)), on a row through the center + r = np.hypot(0.5, np.arange(64) - 31.5) + expected = np.sin(np.pi / 2 * np.clip(1 - r / 32, 0, 1)) ** 2 + assert np.allclose(window_mask(64)[31], expected, atol=1e-6) + + # cosine fringes with g = 0.25 1/A at 60 degrees from the row axis toward the column axis + sampling, k, theta = 0.5, 0.25, np.deg2rad(60.0) + rr, cc = np.mgrid[:256, :256] * sampling + image = 100.0 + 5.0 * np.cos(2 * np.pi * k * (rr * np.cos(theta) + cc * np.sin(theta))) + image += np.random.default_rng(0).normal(0, 1.0, image.shape) + ds = Dataset2d.from_array(image, sampling=[0.05, 0.05], units=["nm", "nm"]) + lf = LocalFFT.from_image(ds, window_size=64, step=32, k_max=0.4) + + assert np.isclose(lf.sampling, sampling), "nm must convert to Angstroms" + dp = lf.dataset.array[3, 3] + h = int(lf.center[0]) + assert dp[h, h] == 0.0, "the center is removed by the normalization" + row, col = np.unravel_index(np.argmax(dp), dp.shape) + q = np.array([row - h, col - h]) * lf.pixel_size + assert np.isclose(np.hypot(*q), k, atol=lf.pixel_size) + assert np.isclose(np.rad2deg(np.arctan2(q[1], q[0])) % 180, 60.0, atol=5.0) + assert 0.5 < np.median(dp[dp != 0] + 1.0) < 1.5, "normalized noise must be close to 1"