diff --git a/README.md b/README.md index 652ba27..6d9038e 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,8 @@ subtraction...` (`Ctrl+Shift+B`). **Performing an FFT.** `Measurements → FFT viewer...` (`Ctrl+Shift+F`) shows the spectrum with q-axes in nm⁻¹, intensity controls, and a radial profile; the tabs fit a reciprocal lattice, correct drift distortion, -suppress mains pickup, and reconstruct a filtered image by inverse FFT. +suppress mains pickup, reconstruct a filtered image by inverse FFT, and +symmetrize the image by rotational averaging. **Finding features.** `Measurements → Feature finder...` detects maxima or minima with threshold, spacing, and smoothing controls, then exports the @@ -100,9 +101,11 @@ suite. What it does today: - **FFT tools** (the FFT viewer) — inspect the magnitude and radial profile with q in nm⁻¹; overlay a draggable reciprocal-lattice grid and apply an affine lattice correction; show Bragg-shell rings for a known structure; predict and - notch out **mains pickup** (50/60 Hz); and use the **inverse-FFT / + notch out **mains pickup** (50/60 Hz); use the **inverse-FFT / Fourier-reconstruction** tool to select circle/ellipse features, *remove* or - *keep* them, preview the reconstructed image and the residual, then apply. + *keep* them, preview the reconstructed image and the residual, then apply; + and **symmetrize** an image by n-fold (optionally mirrored) rotational + averaging, with the removed residual always shown alongside. - **ROIs and measurements** — rectangle, ellipse, polygon, freehand, line, and point ROIs; ROI-scoped processing; line profiles, periodicity, ROI statistics, step heights, distances and angles, feature points, point-mask @@ -112,10 +115,17 @@ suite. What it does today: molecule segmentation, counting, few-shot classification, template-match counting, lattice extraction, and reproducible step-edge exclusion. Requires the `features` extra (OpenCV + scikit-learn). +- **Particle Statistics** *(optional)* — test detected point patterns against + spatial null models (random, hard-core, measured-feature) with simulation + envelopes, plus a guided tutorial and free-play simulations. Model runs + require the `adstat` extra; treat verdicts as exploratory (see the + [guide](docs/adstat_user_guide.md)). - **Spectroscopy** — inspect single traces or overlays / waterfalls. Smoothing, derivative, normalization, outlier masking, and offsets operate on derived display data; the raw loaded arrays are left intact. -- **Convert** Createc `.dat` scans to Nanonis-compatible `.sxm` (and PNG). +- **Convert** Createc `.dat` scans to Nanonis-compatible `.sxm` (and PNG), or + export them as raw / physical NumPy `.npy` bundles with header and + provenance sidecars (`dat2npy`). - **Export with context** — PNG, PDF, CSV, JSON, SXM, and optional Gwyddion `.gwy`. Many exports also write a JSON provenance sidecar (source file and channel, display settings, processing state, ROIs, and warnings). It is not a @@ -133,6 +143,7 @@ suite. What it does today: | Input | Nanonis `.dat` | Point spectroscopy | | Input | RHK `.sm4` | STM/SPM image scan | | Output | `.sxm` | Converted or processed scan data | +| Output | `.npy` | Raw / physical NumPy array bundles (with header + provenance sidecars) | | Output | `.png`, `.pdf` | Figure / image export | | Output | `.csv`, `.json` | Numerical data, metadata, or provenance | | Output | `.gwy` | Optional Gwyddion export when `gwyfile` is installed | @@ -143,18 +154,23 @@ are in [docs/createc_dat_reader.md](docs/createc_dat_reader.md). ## Installation notes Python 3.11 or newer is required. The core install pulls in numpy, scipy, -Pillow, PySide6, matplotlib, and shapely. +Pillow, PySide6, matplotlib, shapely, and scikit-image. Optional extras: ```bash python -m pip install -e ".[features]" # particle/feature + lattice tools (OpenCV, scikit-learn) +python -m pip install -e ".[adstat]" # Particle Statistics engine (adstat) +python -m pip install -e ".[clip]" # CLIP encoder for few-shot classification (torch) python -m pip install -e ".[gwyddion]" # Gwyddion .gwy writer (gwyfile) -python -m pip install -e ".[dev]" # test + lint tooling +python -m pip install -e ".[dev]" # test tooling (pytest, vulture) +python -m pip install -e ".[all]" # features + adstat + gwyddion + pytest (no CLIP) ``` The Feature Counting and lattice-extraction tools are inactive until the -`features` extra is installed; everything else works with the core install. +`features` extra is installed, Particle Statistics model comparisons need the +`adstat` extra, and the CLIP feature encoder needs the `clip` extra; +everything else works with the core install. ## Using ProbeFlow from Python @@ -195,7 +211,7 @@ dzdv = numeric_derivative(spec.x_array, z_smooth) ```bash python -m pip install -e ".[dev,features]" pytest # run the test suite -ruff check probeflow tests # lint +ruff check probeflow tests # lint (install ruff separately or via pre-commit) ``` Repository layout: diff --git a/docs/createc_dat_reader.md b/docs/createc_dat_reader.md index 4e501d2..af81e77 100644 --- a/docs/createc_dat_reader.md +++ b/docs/createc_dat_reader.md @@ -82,7 +82,11 @@ ProbeFlow interprets it as the one-based next Y position, so completed rows are `ImageYPosMax - 1`. If that value is inside the declared image height, the scan is marked partial and decoded arrays are trimmed to the completed rows. -If `ImageYPosMax` is absent or outside the declared height, ProbeFlow falls back -to the existing channel-0 row heuristic. `is_partial_scan`, `image_y_pos_max`, +If `ImageYPosMax` indicates that every declared row completed +(`ImageYPosMax - 1 >= Num.Y`), the stack is returned untrimmed and the +channel-0 row heuristic is skipped — the heuristic could otherwise silently +drop genuine trailing rows whose DAC values happen to be zero. Only when +`ImageYPosMax` is absent (or nonsensical, below 1) does ProbeFlow fall back to +the channel-0 row heuristic. `is_partial_scan`, `image_y_pos_max`, `original_Ny`, and `trimmed_Ny` are recorded in the decode report so callers can distinguish a completed smaller image from an interrupted acquisition. diff --git a/docs/gui.md b/docs/gui.md index b8c6dd2..1d3fc73 100644 --- a/docs/gui.md +++ b/docs/gui.md @@ -84,6 +84,12 @@ axes. The tabs below cover the common reciprocal-space tasks: * **Mains** — detect and suppress mains-frequency pickup streaks. * **Inverse FFT** — mask regions of the spectrum and reconstruct the filtered image. +* **Symmetrize** — enforce an n-fold (optionally mirrored) symmetry by + averaging the image with its rotated copies. Rotated copies are + registered back onto the original automatically, so the symmetry axis + need not sit at the image centre. Always check the **Residual** + preview: it holds everything symmetrization removed — noise, but also + real defects and domain boundaries. **Focus FFT** and the zoom buttons home in on the spectral content near the origin, and the **Export** menu saves the spectrum or filtered image. @@ -135,9 +141,9 @@ and the developer-facing [AdStat integration](adstat_integration.md) contract. ## Beyond the basics -* **ROIs** — draw rectangles, ellipses, and lines from the ROI tab; ROIs - restrict background fits, FFTs, and statistics, and are saved as - sidecar files next to the scan. +* **ROIs** — draw rectangles, ellipses, polygons, freehand outlines, + lines, and points from the ROI tab; ROIs restrict background fits, + FFTs, and statistics, and are saved as sidecar files next to the scan. * **Measurements** — distance and angle measurements, line profiles, ROI statistics, step heights, pair correlation. * **Spectroscopy** — `.VERT` and Nanonis spectroscopy files open in a diff --git a/docs/review_status.md b/docs/review_status.md index b3898ff..4f17766 100644 --- a/docs/review_status.md +++ b/docs/review_status.md @@ -1,6 +1,6 @@ # ProbeFlow Review Status -**Updated**: 2026-06-12 +**Updated**: 2026-07-03 This is the single, consolidated record of ProbeFlow's code-review history. The detailed per-angle review files (`docs/reviews/2026-05-27-*.md`) and the earlier @@ -201,6 +201,102 @@ the per-finding detail. Suite grew from ~2,260 to 2,449 tests across the campaign; main is green. +## Scheduled deep review, 2026-07-03 (new-code pass: .dat→.npy converter, feature bank, CLIP classify, contrast fix) + +Targets: everything merged since the adversarial campaign that earlier passes +had not covered — the Createc `.dat`→`.npy` converter (+ the #46 scan-range +fix), the feature bank Phases 1–2, the CLIP classify encoder, and the +viewer-black contrast fix (8a2950d). Particle Statistics (#45) was itself a +review and was not re-reviewed. + +### Fixed in this pass (uncommitted, suite green: 2508 passed / 3 skipped) + +1. **Feature-bank silent data loss (the one real find)** — + `analysis/feature_bank.py`: a corrupt/unparseable bank file made + `load_bank` silently return an empty bank, so the next "Add to bank" + overwrote the user's entire cross-scan labelled-sample collection with just + the new entries. `append_entries` now loads strict (`ValueError`, surfaced + on the GUI status bar, original file untouched) and writes atomically + (tmp + `os.replace`), closing the truncated-write corruption source too. + Read path (classification) stays permissive. New `tests/test_feature_bank.py` + (the module previously had **zero** tests). +2. **Createc complete-scan trim hazard** — `io/readers/createc_dat.py` + `_trim_createc_stack`: when `ImageYPosMax` confirmed the scan completed + (`= Num.Y + 1`), the code still fell through to the legacy channel-0 + nonzero heuristic, which silently drops genuine trailing rows whose DAC + values are zero (and drops the last row if the bottom-right pixel is 0/NaN). + Header-confirmed-complete now returns untrimmed. Pinned in + `tests/test_createc_dat_decode.py`. +3. **Converter `main()` partial-kwarg crash** — + `io/converters/createc_dat_to_npy.py`: programmatic calls that passed only + some keywords (e.g. `basis=` without `src=`) skipped the CLI arg-parse + branch and crashed on `Path(None)`. Defaults now resolved unconditionally. +4. **Mojibake log line** — `gui/workers.py` NumPy-conversion log printed + `──` (double-encoded box-drawing chars) in the convert dialog. + +### Verified sound (no change needed) + +- `decoded_scan_range_m` (#46): per-pixel step from the programmed frame × + decoded pixel count — physically correct for first-column removal and + partial-scan trims; metadata path consistent (`shape` (Ny, Nx) + decoded range). +- CLIP classify + bank integration (`analysis/features.py`): bank-dimension + guard, zero-norm cosine guard, bank-only degenerate sample stack, sharpness + z-norm padding for bank entries, encoder gating in the GUI controller. +- Contrast revalidation (8a2950d): sound; known tradeoff — a deliberate manual + window capturing <2% of finite data is reset by the next processing op. +- Previously-flagged gaps now confirmed fixed: shear ROI/mask invalidation + (`_on_shear_applied`), quick-selection menu sync after transform-drop. + +### Larger items — outlined fixes (not done; ordered by value) + +1. **Feature-bank scale blindness** (tracked TODO in `feature_bank.py`): crops + are fixed 48 px, so a banked embedding's physical field of view depends on + the source scan resolution and cross-scan matches are scale-blind. + *Plan:* bump `BANK_SCHEMA_VERSION` → 2; `make_entry` gains + `pixel_size_m` (or crop physical size in nm) threaded from the features + panel's scan; `bank_to_samples` returns it; `classify_particles` warns when + bank-entry crop size differs from the current crop's physical size by more + than ~2× (later: resample crops to a canonical nm/px before embedding). + Old v1 entries: treat as unknown scale, warn once. +2. **Background-dominated CLIP embedding** (carried from UniMR review): + similarities crush to ~0.998–1.0 because the 48 px crop is mostly flat + background, so outlier rejection is unreliable in general. *Plan:* mask or + weight the crop by the particle contour before embedding (contours already + exist on `Particle`); re-verify `_threshold_similarities` guards afterwards + — the spread guard may become unnecessary. +3. **Bundle-provenance pixel-key convention**: in the `.npy` bundle sidecars, + `original_shape`/`decoded_shape` are `(Nx, Ny)` while each plane's `shape` + is `(rows, cols)` — same word, opposite conventions in one JSON (e.g. + A250407: `decoded_shape: [255, 9]` vs plane `shape: [9, 255]`). *Plan:* + rename the bundle keys to `original_pixels`/`decoded_pixels` (matching + `scan_pixels`) while the format is young; regenerate the committed sample + bundles under `test_data/output_raw_npy/`; grep for external consumers first. +4. **Bank rotation-augmentation asymmetry**: in-image samples get 36 rotated + copies under `rotate_augment`; bank vectors are single unrotated embeddings, + so banked classes are disadvantaged for rotated instances. *Plan:* either + bank all 36 rotated embeddings per sample (36× file size) or store the raw + crop in the entry and re-embed with rotations at classify time (slower, + smaller file). Decide with real usage data. + +### Next-run candidates (unexplored, carried forward) + +- `processed_export.py` / provenance replay **without** canonical sidecars, and + a CLI replay smoke test of new-format sidecars. +- `image_canvas.py` selection internals (~1500 lines, only ever grepped). +- Spectroscopy display wiring (`gui/spec_viewer/`), `io/readers/createc_vert.py` + + `nanonis_spec.py` decoders, `analysis/spec_plot.py`. +- Partial `.sxm`/`.dat` thumbnail decode in browse (open since the browse + loading review). +- Feature-bank GUI flows under Qt (add-to-bank dialog, `refresh_bank_status`) + have no GUI-level tests. + +Environment note for future runs: the full suite needs +`QT_QPA_PLATFORM_PLUGIN_PATH=$(python -c "import PySide6,pathlib; print(pathlib.Path(PySide6.__file__).parent/'Qt'/'plugins'/'platforms')")` +alongside `QT_QPA_PLATFORM=offscreen` in sandboxed shells, or `qapp` creation +aborts (SIGABRT) at `test_feature_finder.py`; `test_fft_viewer_utils.py` and +`test_lattice_grid.py` still segfault under offscreen Qt on anaconda py3.13 +(environmental, pre-existing). + ## Deferred (not in code-review scope) - A true Python 3.11/3.12 test matrix (only the local interpreter was available). diff --git a/probeflow/analysis/feature_bank.py b/probeflow/analysis/feature_bank.py index 7380d57..361bf30 100644 --- a/probeflow/analysis/feature_bank.py +++ b/probeflow/analysis/feature_bank.py @@ -23,6 +23,7 @@ from __future__ import annotations import json +import os from datetime import datetime, timezone from pathlib import Path from typing import Sequence @@ -41,16 +42,33 @@ def _empty_bank() -> dict: return {"schema_version": BANK_SCHEMA_VERSION, "encoder": "clip", "entries": []} -def load_bank(path) -> dict: - """Load a bank file, or an empty bank when it is missing / unreadable.""" +def load_bank(path, *, strict: bool = False) -> dict: + """Load a bank file, or an empty bank when it is missing / unreadable. + + ``strict=True`` raises :class:`ValueError` instead of returning an empty + bank when the file exists but cannot be parsed as a bank. Read-only + consumers (classification) want the permissive default; writers must use + strict so a corrupt file is never silently replaced, destroying every + previously banked sample. + """ p = Path(path) if not p.exists(): return _empty_bank() try: data = json.loads(p.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError, ValueError): + except (json.JSONDecodeError, OSError, ValueError) as exc: + if strict: + raise ValueError( + f"feature bank {p} exists but is not readable as JSON ({exc}); " + "refusing to overwrite it — repair or remove the file first" + ) from exc return _empty_bank() if not isinstance(data, dict): + if strict: + raise ValueError( + f"feature bank {p} does not contain a JSON object; " + "refusing to overwrite it — repair or remove the file first" + ) return _empty_bank() data.setdefault("schema_version", BANK_SCHEMA_VERSION) data.setdefault("encoder", "clip") @@ -106,7 +124,9 @@ def append_entries(path, new_entries: Sequence[dict]) -> dict: # Materialise up front: a generator would be consumed by the loop below, # making the ``skipped`` count in the summary wrong (negative). new_entries = list(new_entries) - bank = load_bank(p) + # strict: a corrupt existing bank must abort the append, not be silently + # replaced by an empty bank plus the new entries. + bank = load_bank(p, strict=True) seen = { (e.get("source_path"), e.get("particle_index")) for e in bank["entries"] @@ -120,7 +140,11 @@ def append_entries(path, new_entries: Sequence[dict]) -> dict: seen.add(key) added += 1 p.parent.mkdir(parents=True, exist_ok=True) - p.write_text(json.dumps(bank, indent=2), encoding="utf-8") + # Atomic replace: an interrupted write must never leave a truncated bank + # in place of the previous good one. + tmp = p.with_name(p.name + ".tmp") + tmp.write_text(json.dumps(bank, indent=2), encoding="utf-8") + os.replace(tmp, p) return { "added": added, "skipped": len(new_entries) - added, diff --git a/probeflow/core/processing_state.py b/probeflow/core/processing_state.py index 1fc0d0f..9e52393 100644 --- a/probeflow/core/processing_state.py +++ b/probeflow/core/processing_state.py @@ -36,6 +36,7 @@ "periodic_notch_filter", "mains_pickup_suppression", "inverse_fft_filter", + "symmetrize_fft", "linear_undistort", "affine_lattice_correction", "arithmetic", diff --git a/probeflow/gui/dialogs/fft_viewer.py b/probeflow/gui/dialogs/fft_viewer.py index d163443..9c761c5 100644 --- a/probeflow/gui/dialogs/fft_viewer.py +++ b/probeflow/gui/dialogs/fft_viewer.py @@ -22,6 +22,7 @@ from probeflow.gui.dialogs.fft_viewer_lattice_mixin import FFTViewerLatticeMixin from probeflow.gui.dialogs.fft_viewer_mains_mixin import FFTViewerMainsMixin from probeflow.gui.dialogs.fft_viewer_reconstruct_mixin import FFTViewerReconstructMixin +from probeflow.gui.dialogs.fft_viewer_symmetrize_mixin import FFTViewerSymmetrizeMixin from probeflow.gui.no_wheel import install_no_wheel_spinboxes from probeflow.gui.viewer.display_range import DisplayRangeController from probeflow.gui.viewer.display_sliders import DisplaySliderController @@ -83,6 +84,7 @@ class FFTViewerDialog( FFTViewerLatticeMixin, FFTViewerMainsMixin, FFTViewerReconstructMixin, + FFTViewerSymmetrizeMixin, QDialog, ): """Side-by-side real-space / FFT inspection window.""" @@ -171,6 +173,8 @@ def __init__( self._fft_selection_overlay = None self._reconstruct_tab_index = -1 self._reconstruct_preview_active = False + self._symmetrize_tab_index = -1 + self._symmetrize_preview_active = False self._new_image_fn = new_image_fn self._fft_drs = DisplayRangeController(clip_low=0.0, clip_high=100.0, parent=self) @@ -806,6 +810,8 @@ def _build_fft_column(self) -> QVBoxLayout: self._build_mains_tab(), "⚡ Mains") self._reconstruct_tab_index = self._tab_widget.addTab( self._build_reconstruct_tab(), "Inverse FFT") + self._symmetrize_tab_index = self._tab_widget.addTab( + self._build_symmetrize_tab(), "Symmetrize") self._fft_splitter = QSplitter(Qt.Vertical) self._fft_splitter.addWidget(fft_top) diff --git a/probeflow/gui/dialogs/fft_viewer_symmetrize_mixin.py b/probeflow/gui/dialogs/fft_viewer_symmetrize_mixin.py new file mode 100644 index 0000000..f23f4fc --- /dev/null +++ b/probeflow/gui/dialogs/fft_viewer_symmetrize_mixin.py @@ -0,0 +1,424 @@ +"""Symmetrize tab for the FFT viewer — n-fold rotational averaging. + +Mixin split out of ``fft_viewer.py`` (same pattern as the Inverse FFT tab). +It relies on attributes owned by ``FFTViewerDialog`` (``self._arr``, +``self._scan_range_m``, ``self._fft_source`` / ``self._roi_id``, the Bragg +reference controls ``self._bragg_sym_combo`` / ``self._bragg_a_spin`` / +``self._bragg_unit_combo``, the ``self._apply_correction_fn`` / +``self._new_image_fn`` callbacks, ``self._fft_pixel_sizes_m`` and the shared +FFT-preview helpers ``_show_fft_preview`` / ``_hide_fft_preview``). + +Symmetrization fabricates data — it paints the assumed symmetry over defects +and domain boundaries — so the residual (original − result) is presented as a +first-class preview, mirroring the Inverse FFT tab's honesty-first design. +""" + +from __future__ import annotations + +import math + +import numpy as np +from probeflow.gui._tooltips import tip as _tip +from probeflow.gui.typography import ui_font +from PySide6.QtWidgets import ( + QCheckBox, QComboBox, QDoubleSpinBox, QFileDialog, QFrame, QGroupBox, + QHBoxLayout, QLabel, QPushButton, QScrollArea, QSpinBox, QVBoxLayout, + QWidget, +) + + +class FFTViewerSymmetrizeMixin: + """Enforce an n-fold symmetry by rotate-and-average; preview and apply.""" + + def _build_symmetrize_tab(self) -> QWidget: + scroll = QScrollArea() + scroll.setWidgetResizable(True) + scroll.setFrameShape(QFrame.NoFrame) + page = QWidget() + lay = QVBoxLayout(page) + lay.setContentsMargins(8, 8, 8, 8) + lay.setSpacing(6) + + intro = QLabel( + "Enforce an n-fold symmetry by averaging the image with its rotated " + "copies (equivalent to rotating FFT segments onto each other). " + "Always inspect the residual: it is everything symmetrization " + "removed — noise, but also real defects and domain boundaries.") + intro.setWordWrap(True) + intro.setFont(ui_font(9)) + intro.setToolTip(_tip( + "Averaging the image over its symmetry rotations keeps everything " + "n-fold symmetric (the lattice) and suppresses everything that is " + "not (noise, defects) by ~1/n. Each rotated copy is registered " + "back onto the original first, so the symmetry axis does not have " + "to sit at the image centre. The residual = original − result " + "shows exactly what was removed — if a defect matters, it is in " + "there.")) + lay.addWidget(intro) + + # ── symmetry ──────────────────────────────────────────────────────── + sgrp = QGroupBox("Symmetry") + sg = QVBoxLayout(sgrp) + sg.setSpacing(4) + fold_row = QHBoxLayout() + fold_row.addWidget(QLabel("Fold order:")) + self._symm_fold_spin = QSpinBox() + self._symm_fold_spin.setRange(2, 12) + self._symm_fold_spin.setValue(6) + self._symm_fold_spin.setToolTip(_tip( + "Rotation order n: copies rotated by 360°/n are averaged. Use 6 " + "for hexagonal lattices, 4 for square, 3 for threefold surfaces, " + "2 for twofold.")) + self._symm_fold_spin.setMaximumWidth(72) + self._symm_fold_spin.valueChanged.connect( + lambda _v: self._update_symmetrize_status()) + fold_row.addWidget(self._symm_fold_spin) + fold_row.addStretch(1) + sg.addLayout(fold_row) + + self._symm_mirror_cb = QCheckBox("Also mirror (dihedral symmetry)") + self._symm_mirror_cb.setToolTip(_tip( + "Additionally average the mirrored copies (point group D_n instead " + "of C_n). Unlike pure rotations, mirrors need the axis angle below " + "to be aligned with the lattice.")) + self._symm_mirror_cb.toggled.connect(self._on_symm_mirror_toggled) + sg.addWidget(self._symm_mirror_cb) + + axis_row = QHBoxLayout() + self._symm_axis_lbl = QLabel("Mirror axis:") + axis_row.addWidget(self._symm_axis_lbl) + self._symm_axis_spin = QDoubleSpinBox() + self._symm_axis_spin.setRange(0.0, 180.0) + self._symm_axis_spin.setDecimals(1) + self._symm_axis_spin.setSingleStep(0.5) + self._symm_axis_spin.setSuffix(" °") + self._symm_axis_spin.setToolTip(_tip( + "Angle of the mirror line through the image centre. 0° is " + "horizontal (+x). Only used when mirroring is on — pure rotations " + "need no axis.")) + self._symm_axis_spin.setMaximumWidth(90) + axis_row.addWidget(self._symm_axis_spin) + self._symm_align_btn = QPushButton("Align to Bragg peaks") + self._symm_align_btn.setToolTip(_tip( + "Detect the Bragg peaks (using the reference lattice constant and " + "symmetry from the Inspect tab) and set the mirror axis to the " + "measured lattice direction.")) + self._symm_align_btn.clicked.connect(self._on_symm_align_axis) + axis_row.addWidget(self._symm_align_btn) + axis_row.addStretch(1) + sg.addLayout(axis_row) + + # ── options ───────────────────────────────────────────────────────── + ogrp = QGroupBox("Options") + og = QVBoxLayout(ogrp) + og.setSpacing(4) + self._symm_register_cb = QCheckBox("Auto-centre (register rotated copies)") + self._symm_register_cb.setChecked(True) + self._symm_register_cb.setToolTip(_tip( + "Shift each rotated copy back onto the original by cross-" + "correlation before averaging. Keep this on: without it the " + "symmetry axis must pass exactly through the image centre, or the " + "average destroys the lattice contrast instead of cleaning it.")) + og.addWidget(self._symm_register_cb) + interp_row = QHBoxLayout() + interp_row.addWidget(QLabel("Interpolation:")) + self._symm_interp_combo = QComboBox() + self._symm_interp_combo.addItems(["Linear", "Cubic"]) + self._symm_interp_combo.setToolTip(_tip( + "Resampling used for the rotations. Linear is the default; Cubic " + "keeps sharp lattices slightly crisper at the cost of mild " + "ringing.")) + self._symm_interp_combo.setMaximumWidth(100) + interp_row.addWidget(self._symm_interp_combo) + interp_row.addStretch(1) + og.addLayout(interp_row) + self._symm_strict_cb = QCheckBox("Blank pixels not covered by every copy") + self._symm_strict_cb.setToolTip(_tip( + "Rotated copies leave the frame near the corners, so those pixels " + "average fewer copies. Off (default): renormalize by the available " + "copies. On: set them to NaN so the output only contains fully " + "symmetrized pixels.")) + og.addWidget(self._symm_strict_cb) + + top_row = QHBoxLayout() + top_row.setSpacing(6) + top_row.addWidget(sgrp, 1) + top_row.addWidget(ogrp, 1) + lay.addLayout(top_row) + + # ── preview / apply ───────────────────────────────────────────────── + view_row = QHBoxLayout() + view_row.addWidget(QLabel("Preview:")) + self._symm_view_combo = QComboBox() + self._symm_view_combo.addItems(["Result", "Residual", "Original"]) + self._symm_view_combo.setToolTip(_tip( + "Which image the left preview shows: the symmetrized Result, the " + "Residual (everything symmetrization removed), or the Original.")) + self._symm_view_combo.currentIndexChanged.connect( + lambda _=0: self._on_symmetrize_preview() + if self._symmetrize_preview_active else None) + self._symm_view_combo.setMaximumWidth(140) + view_row.addWidget(self._symm_view_combo) + view_row.addStretch(1) + lay.addLayout(view_row) + + pv_row = QHBoxLayout() + self._symm_preview_btn = QPushButton("Preview symmetrized image") + self._symm_preview_btn.setToolTip(_tip( + "Symmetrize with the current settings and show the outcome in the " + "left preview, without changing your data. Check the residual " + "before applying.")) + self._symm_preview_btn.clicked.connect(self._on_symmetrize_preview) + self._symm_clear_btn = QPushButton("Clear preview") + self._symm_clear_btn.setEnabled(False) + self._symm_clear_btn.clicked.connect(self._on_symmetrize_clear) + pv_row.addWidget(self._symm_preview_btn) + pv_row.addWidget(self._symm_clear_btn) + lay.addLayout(pv_row) + + ap_row = QHBoxLayout() + self._symm_apply_btn = QPushButton("Apply to image") + self._symm_apply_btn.setObjectName("accentBtn") + self._symm_apply_btn.setToolTip(_tip( + "Apply the symmetrization. On a whole-image FFT this is recorded " + "as a reproducible step in the processing history. On an ROI it " + "creates a new corrected cropped image (never edits the original " + "in place).")) + self._symm_apply_btn.clicked.connect(self._on_symmetrize_apply) + ap_row.addWidget(self._symm_apply_btn) + lay.addLayout(ap_row) + + ex_row = QHBoxLayout() + exp_res = QPushButton("Export result") + exp_res.setToolTip(_tip("Save the symmetrized image to a file.")) + exp_res.clicked.connect(lambda: self._on_symmetrize_export("result")) + exp_resid = QPushButton("Export residual") + exp_resid.setToolTip(_tip("Save the residual (removed component) to a file.")) + exp_resid.clicked.connect(lambda: self._on_symmetrize_export("residual")) + ex_row.addWidget(exp_res) + ex_row.addWidget(exp_resid) + lay.addLayout(ex_row) + + self._symm_status_lbl = QLabel("") + self._symm_status_lbl.setWordWrap(True) + self._symm_status_lbl.setFont(ui_font(8)) + lay.addWidget(self._symm_status_lbl) + + lay.addStretch(1) + scroll.setWidget(page) + self._on_symm_mirror_toggled(False) + self._update_symmetrize_status() + return scroll + + # ── control state ──────────────────────────────────────────────────────── + + def _on_symm_mirror_toggled(self, checked: bool) -> None: + for w in (self._symm_axis_lbl, self._symm_axis_spin, self._symm_align_btn): + w.setEnabled(bool(checked)) + self._update_symmetrize_status() + + def _symmetrize_kwargs(self) -> dict: + return { + "mirror": bool(self._symm_mirror_cb.isChecked()), + "mirror_axis_deg": float(self._symm_axis_spin.value()), + "register": bool(self._symm_register_cb.isChecked()), + "interpolation": ("linear" + if self._symm_interp_combo.currentIndex() == 0 + else "cubic"), + "strict_coverage": bool(self._symm_strict_cb.isChecked()), + } + + def _anisotropy_warning(self) -> str | None: + px = self._fft_pixel_sizes_m() + if px is None: + return None + px_x, px_y = px + if px_x <= 0 or px_y <= 0: + return None + ratio = max(px_x, px_y) / min(px_x, px_y) + if ratio <= 1.01: + return None + return (f"⚠ Pixels are anisotropic (x/y ratio {ratio:.2f}): rotating " + "in pixel space is not a physical rotation, so the enforced " + "symmetry will be distorted. Resample to square pixels first.") + + # ── align to Bragg ─────────────────────────────────────────────────────── + + def _on_symm_align_axis(self) -> None: + """Set the mirror axis from detected Bragg peaks. + + Uses the reference lattice (constant + symmetry) from the Inspect tab + to predict the first-order Bragg radius, finds the peaks in that + annulus, and takes their common n-fold axis. + """ + arr = np.asarray(self._arr, dtype=np.float64) + from probeflow.processing.bragg import ( + find_bragg_peaks_in_annulus, predicted_bragg_radius) + from probeflow.processing.symmetrize import fold_axis_from_peaks + + symmetry = "square" if self._bragg_sym_combo.currentIndex() == 0 else "hex" + a_val = float(self._bragg_a_spin.value()) + unit = self._bragg_unit_combo.currentText() + a_m = a_val * 1e-10 if unit == "Å" else a_val * 1e-9 + try: + w_m = float(self._scan_range_m[0]) + h_m = float(self._scan_range_m[1]) + scan_m = math.sqrt(w_m * h_m) # geometric mean for non-square scans + r_px = predicted_bragg_radius( + a_m, symmetry, scan_m, min(arr.shape[:2])) + except (ValueError, TypeError) as exc: + self._symm_status_lbl.setText(f"Cannot predict Bragg radius: {exc}") + return + + n = int(self._symm_fold_spin.value()) + finite = np.isfinite(arr) + filled = np.where(finite, arr, float(arr[finite].mean()) if finite.any() else 0.0) + mag = np.abs(np.fft.fftshift(np.fft.fft2(filled - filled.mean()))) + peaks = find_bragg_peaks_in_annulus(mag, r_px, expected_count=n) + + # The detector returns local maxima even when the annulus holds only + # the noise floor (e.g. a wrong lattice constant), so gate the picks on + # how far they stand above the annulus statistics before trusting them + # (real Bragg peaks measure z ≈ 7–10, leakage maxima z ≈ 2–4). + if peaks.shape[0]: + Ny, Nx = mag.shape + cy, cx = Ny / 2.0, Nx / 2.0 + gy, gx = np.ogrid[:Ny, :Nx] + annulus = mag[(np.hypot(gx - cx, gy - cy) >= r_px * 0.8) + & (np.hypot(gx - cx, gy - cy) <= r_px * 1.2)] + mean_a = float(annulus.mean()) + std_a = float(annulus.std()) + 1e-12 + strong = [ + (dx, dy) for dx, dy in peaks + if (mag[int(round(cy + dy)), int(round(cx + dx))] - mean_a) + / std_a >= 5.0] + peaks = np.array(strong, dtype=np.float64).reshape(-1, 2) + if peaks.shape[0] < 2: + self._symm_status_lbl.setText( + "Could not detect clear Bragg peaks near the predicted " + "radius — check the reference lattice constant and symmetry " + "on the Inspect tab, then try again.") + return + try: + axis = fold_axis_from_peaks(peaks, n) + except ValueError: + self._symm_status_lbl.setText( + f"The {peaks.shape[0]} detected peaks do not share a " + f"{n}-fold axis — is the fold order right for this lattice?") + return + self._symm_axis_spin.setValue(axis) + self._symm_status_lbl.setText( + f"Mirror axis set to {axis:.1f}° from {peaks.shape[0]} Bragg " + f"peak{'s' if peaks.shape[0] != 1 else ''} " + f"(axis repeats every {360.0 / n:.0f}°).") + + # ── preview / apply / export ───────────────────────────────────────────── + + def _compute_symmetrization(self, array): + from probeflow.processing.symmetrize import symmetrize_image + return symmetrize_image( + np.asarray(array, dtype=np.float64), + int(self._symm_fold_spin.value()), + **self._symmetrize_kwargs()) + + def _on_symmetrize_preview(self) -> None: + arr = np.asarray(self._arr, dtype=np.float64) + try: + res = self._compute_symmetrization(arr) + except ValueError as exc: + self._symm_status_lbl.setText(f"Symmetrize failed: {exc}") + return + view = self._symm_view_combo.currentText() + disp = {"Result": res.result, "Residual": res.residual, "Original": arr}[view] + self._show_fft_preview(disp) + self._symmetrize_preview_active = True + self._symm_clear_btn.setEnabled(True) + self._update_symmetrize_status(res) + + def _on_symmetrize_clear(self) -> None: + self._hide_fft_preview() + self._symmetrize_preview_active = False + self._symm_clear_btn.setEnabled(False) + self._update_symmetrize_status() + + def _symmetrize_op_params(self) -> dict: + params = {"n_fold": int(self._symm_fold_spin.value()), + **self._symmetrize_kwargs(), + "fft_source": self._fft_source} + if self._fft_source == "active_roi" and self._roi_id is not None: + params["fft_roi_id"] = self._roi_id + return params + + def _on_symmetrize_apply(self) -> None: + params = self._symmetrize_op_params() + if self._fft_source == "active_roi": + self._apply_symmetrize_to_new_image(params) + return + if self._apply_correction_fn is None: + self._symm_status_lbl.setText("Apply is unavailable in this context.") + return + if self._symmetrize_preview_active: + self._hide_fft_preview() + self._symmetrize_preview_active = False + self._symm_clear_btn.setEnabled(False) + self._apply_correction_fn("symmetrize_fft", params) + self._symm_status_lbl.setText( + f"Applied {params['n_fold']}-fold symmetrization.") + + def _apply_symmetrize_to_new_image(self, params: dict) -> None: + arr = np.asarray(self._arr, dtype=np.float64) + try: + res = self._compute_symmetrization(arr) + except ValueError as exc: + self._symm_status_lbl.setText(f"Symmetrize failed: {exc}") + return + fn = getattr(self, "_new_image_fn", None) + if fn is None: + self._symm_status_lbl.setText( + "ROI apply needs a host that can open a new image — export the " + "result instead.") + return + fn(res.result, tuple(self._scan_range_m), + {"op": "symmetrize_fft", "params": params}) + self._symm_status_lbl.setText("Created a new symmetrized image from the ROI.") + + def _on_symmetrize_export(self, which: str) -> None: + arr = np.asarray(self._arr, dtype=np.float64) + try: + res = self._compute_symmetrization(arr) + except ValueError as exc: + self._symm_status_lbl.setText(f"Symmetrize failed: {exc}") + return + data = res.result if which == "result" else res.residual + path, _ = QFileDialog.getSaveFileName( + self, f"Export {which}", f"symmetrized_{which}.csv", + "CSV (*.csv);;NumPy (*.npy)") + if not path: + return + try: + if path.lower().endswith(".npy"): + np.save(path, data) + else: + np.savetxt(path, np.nan_to_num(data), delimiter=",") + self._symm_status_lbl.setText(f"Exported {which} → {path}") + except Exception as exc: + self._symm_status_lbl.setText(f"Export failed: {exc}") + + def _update_symmetrize_status(self, res=None) -> None: + if not getattr(self, "_symm_status_lbl", None): + return + n = int(self._symm_fold_spin.value()) + group = f"D{n}" if self._symm_mirror_cb.isChecked() else f"C{n}" + src = "ROI" if self._fft_source == "active_roi" else "whole image" + bits = [f"Symmetry: {group} ({n}-fold" + f"{' + mirror' if self._symm_mirror_cb.isChecked() else ''}) · " + f"FFT source: {src}"] + if res is not None: + pct = res.symmetry_residual_norm * 100.0 + bits.append( + f"Removed asymmetry: {pct:.1f}% of image contrast " + "(inspect the Residual preview — defects live there)") + warn = self._anisotropy_warning() + if warn: + bits.append(warn) + self._symm_status_lbl.setText("\n".join(bits)) diff --git a/probeflow/gui/workers.py b/probeflow/gui/workers.py index 195d79d..b728751 100644 --- a/probeflow/gui/workers.py +++ b/probeflow/gui/workers.py @@ -552,7 +552,7 @@ def _log(msg, tag="info"): if self.do_npy_raw or self.do_npy_physical: from probeflow.io.converters.createc_dat_to_npy import main as npy_main - _log("── NumPy conversion ──", "info") + _log("── NumPy conversion ──", "info") basis = ( "both" if self.do_npy_raw and self.do_npy_physical diff --git a/probeflow/io/converters/createc_dat_to_npy.py b/probeflow/io/converters/createc_dat_to_npy.py index e8fd337..9806f39 100644 --- a/probeflow/io/converters/createc_dat_to_npy.py +++ b/probeflow/io/converters/createc_dat_to_npy.py @@ -374,6 +374,13 @@ def main( force = bool(args.force) verbose = bool(args.verbose) + # Programmatic callers may pass only some keywords; fall back to the same + # defaults the CLI path uses instead of crashing on Path(None). + if src is None: + src = DEFAULT_INPUT_DIR + if out_root is None: + out_root = DEFAULT_OUTPUT_DIR + setup_logging(verbose) src_path = Path(src) diff --git a/probeflow/io/readers/createc_dat.py b/probeflow/io/readers/createc_dat.py index df0f99d..79b8986 100644 --- a/probeflow/io/readers/createc_dat.py +++ b/probeflow/io/readers/createc_dat.py @@ -388,6 +388,11 @@ def _trim_createc_stack( rows_from_header = image_y_pos_max - 1 if 1 <= rows_from_header < Ny: return stack[:, :rows_from_header, :], rows_from_header + if rows_from_header >= Ny: + # The header confirms every declared row completed. Skip the + # legacy channel-0 nonzero heuristic, which would silently drop + # genuine trailing rows whose DAC values happen to be zero. + return stack, Ny return trim_stack(stack) diff --git a/probeflow/processing/gui_adapter.py b/probeflow/processing/gui_adapter.py index f1f225f..70db817 100644 --- a/probeflow/processing/gui_adapter.py +++ b/probeflow/processing/gui_adapter.py @@ -60,6 +60,7 @@ _FILTER_OPS_PASSTHROUGH: frozenset[str] = frozenset({ "mains_pickup_suppression", "inverse_fft_filter", + "symmetrize_fft", }) diff --git a/probeflow/processing/image.py b/probeflow/processing/image.py index b5b8960..f606628 100644 --- a/probeflow/processing/image.py +++ b/probeflow/processing/image.py @@ -7,6 +7,7 @@ from probeflow.processing.edge_detection import * # noqa: F401,F403 from probeflow.processing.mains_pickup import * # noqa: F401,F403 from probeflow.processing.inverse_fft import * # noqa: F401,F403 +from probeflow.processing.symmetrize import * # noqa: F401,F403 from probeflow.processing.analysis import * # noqa: F401,F403 from probeflow.processing.png_export import * # noqa: F401,F403 from probeflow.processing.tv import * # noqa: F401,F403 diff --git a/probeflow/processing/state.py b/probeflow/processing/state.py index cbec722..7be3cd4 100644 --- a/probeflow/processing/state.py +++ b/probeflow/processing/state.py @@ -706,6 +706,19 @@ def apply_processing_state( conjugate=bool(p.get("conjugate_symmetric", True)), soft_px=float(p.get("soft_px", 0.0)), ) + elif step.op == "symmetrize_fft": + # Rotate-and-average n-fold symmetrization. Registration shifts + # are recomputed from the data on replay (deterministic), so only + # the symmetry parameters are stored. + a = _proc.symmetrize_filter( + a, + int(p.get("n_fold", 1)), + mirror=bool(p.get("mirror", False)), + mirror_axis_deg=float(p.get("mirror_axis_deg", 0.0)), + register=bool(p.get("register", True)), + interpolation=str(p.get("interpolation", "linear")), + strict_coverage=bool(p.get("strict_coverage", False)), + ) elif step.op == "linear_undistort": a = _proc.linear_undistort( a, diff --git a/probeflow/processing/symmetrize.py b/probeflow/processing/symmetrize.py new file mode 100644 index 0000000..f5b68a7 --- /dev/null +++ b/probeflow/processing/symmetrize.py @@ -0,0 +1,307 @@ +"""Rotational / dihedral symmetrization backend. + +Enforce an n-fold rotational (optionally dihedral, i.e. + mirror) symmetry on +an image by averaging its symmetry-equivalent copies, and return the result +**plus the residual** (``original − result``) so the user can see exactly what +the symmetrization removed. Symmetrization fabricates data — it paints the +assumed symmetry over defects, domain boundaries, and drift — so the residual +is a first-class output, not a diagnostic afterthought (same philosophy as +:mod:`probeflow.processing.inverse_fft`). + +Design notes +------------ +* Rotating the *real-space* image by ``k·360°/n`` and averaging is + mathematically identical to rotating the complex FFT in ``k·360°/n`` + segments about DC and averaging (rotation about the origin commutes with the + Fourier transform; translations only add phase). We do it in real space to + avoid interpolating complex spectra, whose phase wraps make interpolation + artifacts much worse. +* **Registration is load-bearing.** Averaging a lattice with its rotated copy + is constructive only when the rotation axis passes through a symmetry centre + of the pattern; an arbitrary axis translates the rotated lattice by a + non-lattice vector and the average loses contrast (each Fourier component is + scaled by ``(1 + e^{iq·t})/2``). Each rotated copy is therefore registered + back onto the original by FFT cross-correlation (with parabolic sub-pixel + refinement) before averaging. ``register=False`` disables this for the + pre-centred case. +* Rotation leaves the frame at the corners. A validity weight is rotated and + shifted alongside each copy, and the average renormalizes per pixel by the + accumulated weight (``coverage``). ``strict_coverage=True`` instead returns + NaN wherever any copy is missing. Pixels that are NaN in the input carry + zero weight, so symmetry-equivalent data can fill them in. +* Angles follow the ``atan2(dy, dx)`` convention used by + :func:`probeflow.processing.bragg.find_bragg_peaks_in_annulus` peak offsets: + measured from the +x (column) axis toward +y (row) axis. Pixels are assumed + square; on anisotropic pixel scales a pixel-space rotation is not a physical + rotation (the GUI should warn before offering the tool there). + +GUI-free (numpy + scipy.ndimage); import lazily from the FFT viewer / CLI. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import numpy as np +from scipy.ndimage import rotate as _nd_rotate +from scipy.ndimage import shift as _nd_shift + +__all__ = [ + "SymmetrizeResult", + "fold_axis_from_peaks", + "symmetrize_image", + "symmetrize_filter", +] + +# Below this accumulated weight a pixel has no meaningful contribution and the +# renormalized average would amplify interpolation dust; emit NaN instead. +_MIN_WEIGHT = 0.05 + + +@dataclass +class SymmetrizeResult: + result: np.ndarray # symmetrized real-space image (NaN where uncovered) + residual: np.ndarray # original − result (NaN where either is NaN) + symmetrized_fft: np.ndarray # |FFT(result)|, fftshift-centred, for display + coverage: np.ndarray # accumulated validity weight per pixel (max = op count) + shifts: np.ndarray # (n_ops, 2) registration shifts (dy, dx) in px + symmetry_residual_norm: float # rms(residual) / rms(original − mean) — ~0 when symmetric + n_ops: int # symmetry operations averaged (n_fold, ×2 if mirror) + + +def _rotated_copy(filled: np.ndarray, weight: np.ndarray, + angle_deg: float, mirrored: bool, order: int): + """Rotate (and optionally pre-mirror) an image and its validity weight. + + Positive ``angle_deg`` moves a feature at ``atan2(dy, dx)`` angle φ to + φ + angle. ``scipy.ndimage.rotate`` with default ``axes=(1, 0)`` rotates + the opposite way in (row, col) space, hence the sign flip. The mirror is + a row flip (y → −y about the array centre), applied before the rotation so + the pair composes to a reflection across the line at ``angle_deg / 2``. + """ + img, w = filled, weight + if mirrored: + img, w = np.flipud(img), np.flipud(w) + if angle_deg % 360.0 == 0.0: + return img.copy(), w.copy() + kw = dict(reshape=False, order=order, mode="constant", cval=0.0, + prefilter=(order > 1)) + return (_nd_rotate(img, -angle_deg, **kw), + np.clip(_nd_rotate(w, -angle_deg, **kw), 0.0, 1.0)) + + +def _register_shift(ref_fft: np.ndarray, copy_img: np.ndarray, + copy_w: np.ndarray, mean_val: float, + max_shift_px: int) -> tuple[float, float]: + """Best (dy, dx) translation aligning ``copy`` onto the reference. + + FFT cross-correlation of the demeaned, validity-weighted images, peak + search restricted to ``±max_shift_px`` around zero shift (the registration + shift is only defined modulo a lattice vector — we want the branch nearest + zero), 3-point parabolic sub-pixel refinement per axis. + """ + b_c = (copy_img - mean_val) * copy_w + # ref_fft is fft2 of the demeaned, weighted reference; the correlation + # c[s] = Σ ref(x)·copy(x − s) peaks at the shift to apply to the copy. + cc = np.fft.ifft2(ref_fft * np.conj(np.fft.fft2(b_c))).real + cc = np.fft.fftshift(cc) + Ny, Nx = cc.shape + cy, cx = Ny // 2, Nx // 2 + + r = int(max_shift_px) + win = cc[max(0, cy - r): cy + r + 1, max(0, cx - r): cx + r + 1] + iy, ix = np.unravel_index(int(np.argmax(win)), win.shape) + py, px = iy + max(0, cy - r), ix + max(0, cx - r) + + def _parabolic(v_m, v_0, v_p) -> float: + denom = v_m - 2.0 * v_0 + v_p + if abs(denom) < 1e-30: + return 0.0 + return float(np.clip(0.5 * (v_m - v_p) / denom, -0.5, 0.5)) + + dy = float(py - cy) + dx = float(px - cx) + if 0 < py < Ny - 1: + dy += _parabolic(cc[py - 1, px], cc[py, px], cc[py + 1, px]) + if 0 < px < Nx - 1: + dx += _parabolic(cc[py, px - 1], cc[py, px], cc[py, px + 1]) + return dy, dx + + +def symmetrize_image( + image: np.ndarray, + n_fold: int, + *, + mirror: bool = False, + mirror_axis_deg: float = 0.0, + register: bool = True, + max_shift_px: int | None = None, + interpolation: str = "linear", + strict_coverage: bool = False, +) -> SymmetrizeResult: + """Average an image over the C_n (or D_n, with ``mirror``) symmetry group. + + Parameters + ---------- + image + 2-D real-space image. NaNs are treated as missing data (zero weight) + and may be filled in by symmetry-equivalent pixels. + n_fold + Rotation order ``n ≥ 1``; copies are rotated by ``k·360°/n``. + ``n_fold=1`` with ``mirror=False`` is the identity. + mirror + Also average the ``n`` mirrored copies (dihedral group D_n). + mirror_axis_deg + Angle of the mirror line through the image centre, ``atan2(dy, dx)`` + convention. Irrelevant for pure rotations; for mirrors, align it to + the lattice (see :func:`fold_axis_from_peaks`). + register + Register every copy back onto the original by cross-correlation before + averaging. Without this, a symmetry axis that does not pass through + the image centre destroys contrast instead of denoising. + max_shift_px + Half-width of the registration search window (default + ``min(shape) // 4``). The optimum is only defined modulo a lattice + vector, so the window keeps the branch nearest zero shift. + interpolation + ``"linear"`` (default, order 1 — matches the geometry ops) or + ``"cubic"`` (order 3, sharper lattices at slightly more ringing). + strict_coverage + When True, output NaN wherever any symmetry copy left the frame + (coverage below the full operation count) instead of renormalizing. + """ + a = np.asarray(image, dtype=np.float64) + if a.ndim != 2: + raise ValueError("symmetrize_image expects a 2-D image") + n = int(n_fold) + if n < 1: + raise ValueError("n_fold must be >= 1") + if interpolation not in ("linear", "cubic"): + raise ValueError("interpolation must be 'linear' or 'cubic'") + order = 1 if interpolation == "linear" else 3 + + valid = np.isfinite(a) + if not valid.any(): + raise ValueError("image contains no finite pixels") + mean_val = float(a[valid].mean()) + filled = np.where(valid, a, mean_val) + w0 = valid.astype(np.float64) + + if max_shift_px is None: + max_shift_px = max(1, min(a.shape) // 4) + + # (angle, mirrored) members of C_n / D_n. A mirror across the line at + # angle φ is R(2φ) ∘ (y → −y); composing with the k-th rotation gives + # R(k·α + 2φ) ∘ flip. + alpha = 360.0 / n + ops: list[tuple[float, bool]] = [(k * alpha, False) for k in range(n)] + if mirror: + ops += [(k * alpha + 2.0 * float(mirror_axis_deg), True) for k in range(n)] + n_ops = len(ops) + + if n_ops == 1: + result = np.where(valid, a, np.nan) + return SymmetrizeResult( + result=result, residual=np.where(valid, 0.0, np.nan), coverage=w0, + symmetrized_fft=np.abs(np.fft.fftshift(np.fft.fft2(filled))), + shifts=np.zeros((1, 2)), symmetry_residual_norm=0.0, n_ops=1) + + # Skip registration on (near-)constant images: the correlation surface is + # flat and argmax would return an arbitrary corner of the window. + contrast = float(a[valid].std()) + do_register = register and contrast > 1e-12 * (abs(mean_val) + 1.0) + ref_fft = np.fft.fft2((filled - mean_val) * w0) if do_register else None + + num = np.zeros_like(a) + den = np.zeros_like(a) + shifts = np.zeros((n_ops, 2), dtype=np.float64) + for i, (angle, mirrored) in enumerate(ops): + if not mirrored and angle % 360.0 == 0.0: + num += filled * w0 + den += w0 + continue + img, w = _rotated_copy(filled, w0, angle, mirrored, order) + if do_register: + dy, dx = _register_shift(ref_fft, img, w, mean_val, max_shift_px) + if dy != 0.0 or dx != 0.0: + skw = dict(order=1, mode="constant", cval=0.0) + img = _nd_shift(img, (dy, dx), **skw) + w = np.clip(_nd_shift(w, (dy, dx), **skw), 0.0, 1.0) + shifts[i] = (dy, dx) + num += img * w + den += w + + with np.errstate(invalid="ignore", divide="ignore"): + result = np.where(den > _MIN_WEIGHT, num / np.maximum(den, 1e-12), np.nan) + if strict_coverage: + result = np.where(den >= n_ops - 1e-3, result, np.nan) + + residual = np.where(valid & np.isfinite(result), a - result, np.nan) + res_valid = np.isfinite(residual) + if res_valid.any(): + rms_res = float(np.sqrt(np.nanmean(residual[res_valid] ** 2))) + rms_sig = float(np.sqrt(np.mean((a[valid] - mean_val) ** 2))) + symmetry_residual_norm = rms_res / (rms_sig + 1e-12) + else: + symmetry_residual_norm = float("nan") + + fft_src = np.where(np.isfinite(result), result, mean_val) + return SymmetrizeResult( + result=result, + residual=residual, + symmetrized_fft=np.abs(np.fft.fftshift(np.fft.fft2(fft_src))), + coverage=den, + shifts=shifts, + symmetry_residual_norm=symmetry_residual_norm, + n_ops=n_ops, + ) + + +def symmetrize_filter( + image: np.ndarray, + n_fold: int, + *, + mirror: bool = False, + mirror_axis_deg: float = 0.0, + register: bool = True, + interpolation: str = "linear", + strict_coverage: bool = False, +) -> np.ndarray: + """Convenience: return only the symmetrized image. + + This is the entry point for the (future) ``symmetrize_fft`` + ProcessingState op, mirroring :func:`probeflow.processing.inverse_fft. + inverse_fft_filter`. ``n_fold=1`` without ``mirror`` returns a copy. + """ + return symmetrize_image( + np.asarray(image, dtype=np.float64), n_fold, + mirror=mirror, mirror_axis_deg=mirror_axis_deg, register=register, + interpolation=interpolation, strict_coverage=strict_coverage, + ).result + + +def fold_axis_from_peaks(peaks_xy: np.ndarray, n_fold: int) -> float: + """Lattice axis angle (degrees, mod ``360/n``) from Bragg peak offsets. + + ``peaks_xy`` is the ``(M, 2)`` array of (x, y) pixel offsets from DC + returned by :func:`probeflow.processing.bragg.find_bragg_peaks_in_annulus`. + Under an n-fold symmetry the peak angles are ``θ0 + k·360/n``; multiplying + by ``n`` collapses them onto one direction, whose circular mean is robust + to a missing or spurious peak. Use the result to align + ``mirror_axis_deg`` (or a GUI overlay) with the lattice. + """ + pts = np.asarray(peaks_xy, dtype=np.float64) + if pts.ndim != 2 or pts.shape[1] != 2 or pts.shape[0] == 0: + raise ValueError("peaks_xy must be a non-empty (M, 2) array") + n = int(n_fold) + if n < 1: + raise ValueError("n_fold must be >= 1") + theta = np.arctan2(pts[:, 1], pts[:, 0]) + # Weight the circular mean by peak radius so a weak spurious pick near DC + # cannot swing the axis. + weights = np.hypot(pts[:, 0], pts[:, 1]) + z = np.sum(weights * np.exp(1j * n * theta)) + if abs(z) < 1e-12: + raise ValueError("peak angles have no consistent n-fold axis") + return math.degrees(np.angle(z) / n) % (360.0 / n) diff --git a/tests/test_createc_dat_decode.py b/tests/test_createc_dat_decode.py index 3348eb1..9c4c829 100644 --- a/tests/test_createc_dat_decode.py +++ b/tests/test_createc_dat_decode.py @@ -430,6 +430,30 @@ def test_image_y_pos_max_records_partial_scan_without_guessing_full_height(tmp_p assert report.decoded_channels_dac.shape == (1, 3, 2) +def test_image_y_pos_max_complete_scan_keeps_zero_valued_trailing_rows(tmp_path): + # ImageYPosMax = Num.Y + 1 means every declared row completed; the legacy + # channel-0 nonzero heuristic must not trim genuine all-zero trailing rows. + dat = tmp_path / "complete_flat_tail.dat" + header = ( + b"[Paramco32]\n" + b"Num.X=3\n" + b"Num.Y=4\n" + b"Channels=1\n" + b"ImageYPosMax=5\n" + ) + data = np.arange(1, 13, dtype="= 0 + assert dlg._tab_widget.tabText(dlg._symmetrize_tab_index) == "Symmetrize" + assert dlg._symm_fold_spin.value() == 6 + assert dlg._symm_register_cb.isChecked() + dlg.close() + + def test_mirror_toggle_gates_axis_controls(self, qapp): + dlg, _ = _dialog(qapp) + assert not dlg._symm_axis_spin.isEnabled() + assert not dlg._symm_align_btn.isEnabled() + dlg._symm_mirror_cb.setChecked(True) + assert dlg._symm_axis_spin.isEnabled() + assert dlg._symm_align_btn.isEnabled() + dlg.close() + + +class TestPreview: + def test_preview_shows_result_and_metric(self, qapp): + dlg, _ = _dialog(qapp) + dlg._on_symmetrize_preview() + assert dlg._symmetrize_preview_active + assert dlg._symm_clear_btn.isEnabled() + assert "Removed asymmetry" in dlg._symm_status_lbl.text() + dlg._on_symmetrize_clear() + assert not dlg._symmetrize_preview_active + dlg.close() + + def test_centred_lattice_reports_low_asymmetry(self, qapp): + dlg, img = _dialog(qapp) + res = dlg._compute_symmetrization(img) + assert res.symmetry_residual_norm < 0.06 + + def test_anisotropic_pixels_warn(self, qapp): + dlg, _ = _dialog(qapp, scan_range=(10e-9, 20e-9)) + assert "anisotropic" in dlg._symm_status_lbl.text() + dlg.close() + + def test_square_pixels_do_not_warn(self, qapp): + dlg, _ = _dialog(qapp) + assert "anisotropic" not in dlg._symm_status_lbl.text() + dlg.close() + + +class TestApply: + def test_apply_emits_symmetrize_op(self, qapp): + captured = {} + dlg, _ = _dialog(qapp, captured) + dlg._symm_fold_spin.setValue(4) + dlg._on_symmetrize_apply() + assert captured["op"] == "symmetrize_fft" + p = captured["params"] + assert p["n_fold"] == 4 + assert p["mirror"] is False + assert p["register"] is True + assert p["interpolation"] == "linear" + assert p["strict_coverage"] is False + assert p["fft_source"] == "whole_image" + dlg.close() + + def test_apply_params_replay_through_processing_state(self, qapp): + # The emitted op must round-trip through the real pipeline. + from probeflow.processing.state import ProcessingState, ProcessingStep, \ + apply_processing_state + from probeflow.processing.symmetrize import symmetrize_filter + captured = {} + dlg, img = _dialog(qapp, captured) + dlg._on_symmetrize_apply() + state = ProcessingState( + steps=[ProcessingStep(captured["op"], captured["params"])]) + out = apply_processing_state(img, state) + np.testing.assert_array_equal(out, symmetrize_filter(img, 6)) + dlg.close() + + def test_apply_hides_active_preview(self, qapp): + captured = {} + dlg, _ = _dialog(qapp, captured) + dlg._on_symmetrize_preview() + dlg._on_symmetrize_apply() + assert not dlg._symmetrize_preview_active + assert captured["op"] == "symmetrize_fft" + dlg.close() + + def test_roi_source_creates_new_image(self, qapp): + opened = {} + + def _new_image(arr, scan_range_m, provenance): + opened.update(arr=arr, scan_range=scan_range_m, provenance=provenance) + + dlg, img = _dialog(qapp, roi=True, new_image_fn=_new_image) + dlg._fft_source = "active_roi" + dlg._on_symmetrize_apply() + assert opened["provenance"]["op"] == "symmetrize_fft" + assert opened["provenance"]["params"]["fft_source"] == "active_roi" + assert opened["provenance"]["params"]["fft_roi_id"] == "roi1" + assert opened["arr"].shape == np.asarray(dlg._arr).shape + dlg.close() + + +class TestAlignToBragg: + def _configured(self, qapp): + dlg, img = _dialog(qapp) + dlg._bragg_sym_combo.setCurrentIndex(1) # Hexagonal + dlg._bragg_unit_combo.setCurrentText("nm") + dlg._bragg_a_spin.setValue(A_NM) + dlg._symm_mirror_cb.setChecked(True) + return dlg + + def test_axis_snaps_to_lattice_direction(self, qapp): + dlg = self._configured(qapp) + dlg._on_symm_align_axis() + assert dlg._symm_axis_spin.value() == pytest.approx(ANGLE0_DEG, abs=2.5) + assert "Mirror axis set" in dlg._symm_status_lbl.text() + dlg.close() + + def test_bad_lattice_constant_reports_no_peaks(self, qapp): + dlg = self._configured(qapp) + dlg._bragg_a_spin.setValue(A_NM * 3.0) # annulus far from real peaks + before = dlg._symm_axis_spin.value() + dlg._on_symm_align_axis() + txt = dlg._symm_status_lbl.text() + assert ("Could not detect" in txt) or ("do not share" in txt) + assert dlg._symm_axis_spin.value() == before + dlg.close() diff --git a/tests/test_pipeline_connectivity.py b/tests/test_pipeline_connectivity.py index 9e62235..5ef494a 100644 --- a/tests/test_pipeline_connectivity.py +++ b/tests/test_pipeline_connectivity.py @@ -55,6 +55,10 @@ "selections": [{"dx": 2, "dy": 0, "rx": 1.0, "ry": 1.0, "angle_deg": 0.0}], "mode": "remove_selected", "conjugate_symmetric": True, "soft_px": 0.0, }, + "symmetrize_fft": { + "n_fold": 2, "mirror": False, "mirror_axis_deg": 0.0, + "register": True, "interpolation": "linear", "strict_coverage": False, + }, "linear_undistort": {"shear_x": 0.0, "scale_y": 1.0}, "affine_lattice_correction": { "matrix": [[1.0, 0.0], [0.0, 1.0]], "expand_canvas": False, diff --git a/tests/test_symmetrize.py b/tests/test_symmetrize.py new file mode 100644 index 0000000..c1e5a0d --- /dev/null +++ b/tests/test_symmetrize.py @@ -0,0 +1,277 @@ +"""Tests for probeflow.processing.symmetrize — n-fold symmetrization backend.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from probeflow.processing.symmetrize import ( + SymmetrizeResult, + fold_axis_from_peaks, + symmetrize_filter, + symmetrize_image, +) + +N = 128 +YY, XX = np.mgrid[:N, :N].astype(np.float64) +# scipy.ndimage.rotate pivots about the (N−1)/2 point; build patterns there so +# "centred" tests exercise the register=False path honestly. +C = (N - 1) / 2.0 + + +def _hex_lattice(center=(C, C), angle0_deg=0.0, period_px=14.0): + """Sum of three cosines with a sixfold reciprocal set (smooth, band-limited).""" + y0, x0 = center + q = 2.0 * np.pi / period_px + out = np.zeros((N, N)) + for k in range(3): + th = np.radians(angle0_deg) + k * np.pi / 3.0 + out += np.cos(q * ((XX - x0) * np.cos(th) + (YY - y0) * np.sin(th))) + return out + + +def _square_lattice(period_px=16.0): + return (np.cos(2.0 * np.pi * (XX - C) / period_px) + + np.cos(2.0 * np.pi * (YY - C) / period_px)) + + +HEX = _hex_lattice() + + +# ─── identity and validation ──────────────────────────────────────────────── + +class TestBasics: + def test_identity_returns_copy(self): + res = symmetrize_image(HEX, 1) + np.testing.assert_array_equal(res.result, HEX) + assert res.n_ops == 1 + assert res.symmetry_residual_norm == 0.0 + assert not np.shares_memory(res.result, HEX) + + def test_rejects_non_2d(self): + with pytest.raises(ValueError, match="2-D"): + symmetrize_image(np.zeros(16), 4) + + def test_rejects_bad_fold(self): + with pytest.raises(ValueError, match="n_fold"): + symmetrize_image(HEX, 0) + + def test_rejects_all_nan(self): + with pytest.raises(ValueError, match="finite"): + symmetrize_image(np.full((8, 8), np.nan), 4) + + def test_rejects_bad_interpolation(self): + with pytest.raises(ValueError, match="interpolation"): + symmetrize_image(HEX, 4, interpolation="nearest") + + def test_filter_matches_result(self): + np.testing.assert_array_equal( + symmetrize_filter(HEX, 6), symmetrize_image(HEX, 6).result) + + def test_result_type(self): + res = symmetrize_image(HEX, 6) + assert isinstance(res, SymmetrizeResult) + assert res.result.shape == HEX.shape + assert res.symmetrized_fft.shape == HEX.shape + assert res.shifts.shape == (6, 2) + assert res.n_ops == 6 + + +# ─── symmetry preservation ────────────────────────────────────────────────── + +class TestPreservation: + def test_c4_on_square_lattice_is_near_exact(self): + # 90° rotations need no interpolation — the average should be exact. + res = symmetrize_image(_square_lattice(), 4) + assert res.symmetry_residual_norm < 1e-10 + + def test_c6_on_centred_hex_lattice(self): + res = symmetrize_image(HEX, 6, register=False) + assert res.symmetry_residual_norm < 0.05 # interpolation loss only + + def test_no_registration_needed_when_centred(self): + res = symmetrize_image(HEX, 6, register=True) + assert np.abs(res.shifts).max() < 0.6 + + def test_odd_sized_image(self): + a = HEX[:127, :127] + res = symmetrize_image(a, 6) + assert res.result.shape == a.shape + assert res.symmetry_residual_norm < 0.06 + + +# ─── registration ─────────────────────────────────────────────────────────── + +class TestRegistration: + """Averaging is constructive only about a symmetry centre; registration + recovers an off-centre axis, which is the tool's default operating mode.""" + + OFF = _hex_lattice(center=(C + 5.3, C - 3.1)) + + def test_registered_off_centre_lattice_survives(self): + res = symmetrize_image(self.OFF, 6, register=True) + assert res.symmetry_residual_norm < 0.06 + + def test_unregistered_off_centre_lattice_degrades(self): + res = symmetrize_image(self.OFF, 6, register=False) + assert res.symmetry_residual_norm > 0.5 + + def test_shifts_are_reported(self): + res = symmetrize_image(self.OFF, 6, register=True) + assert np.abs(res.shifts[1:]).max() > 1.0 + np.testing.assert_array_equal(res.shifts[0], (0.0, 0.0)) + + def test_constant_image_skips_registration(self): + res = symmetrize_image(np.ones((64, 64)), 6, register=True) + np.testing.assert_array_equal(res.shifts, np.zeros((6, 2))) + assert res.symmetry_residual_norm < 1e-10 + + +# ─── denoising / residual honesty ─────────────────────────────────────────── + +class TestResidual: + def test_noise_is_reduced(self): + rng = np.random.default_rng(42) + noisy = HEX + rng.normal(0.0, 0.5, HEX.shape) + res = symmetrize_image(noisy, 6) + full = res.coverage >= res.n_ops - 1e-3 + before = np.std((noisy - HEX)[full]) + after = np.std((res.result - HEX)[full]) + assert after < 0.5 * before + + def test_defect_lands_in_residual(self): + # A sixfold average keeps 1/6 of a lone defect, so the residual at the + # defect carries (n−1)/n of its amplitude. + defect = 3.0 * np.exp(-(((XX - 40) ** 2 + (YY - 80) ** 2) / 30.0)) + res = symmetrize_image(HEX + defect, 6) + assert res.residual[80, 40] > 0.7 * 3.0 * (5.0 / 6.0) + iy, ix = np.unravel_index(np.nanargmax(res.residual), res.residual.shape) + assert abs(iy - 80) <= 2 and abs(ix - 40) <= 2 + + def test_residual_norm_scales_with_asymmetry(self): + sym = symmetrize_image(HEX, 6).symmetry_residual_norm + bumped = HEX + 3.0 * np.exp(-(((XX - 40) ** 2 + (YY - 80) ** 2) / 30.0)) + asym = symmetrize_image(bumped, 6).symmetry_residual_norm + assert asym > 5.0 * sym + + +# ─── coverage and NaN handling ────────────────────────────────────────────── + +class TestCoverage: + def test_corners_have_partial_coverage(self): + res = symmetrize_image(HEX, 6) + assert res.coverage[2, 2] < res.n_ops - 0.5 + assert res.coverage[N // 2, N // 2] == pytest.approx(res.n_ops) + + def test_strict_coverage_nans_corners(self): + res = symmetrize_image(HEX, 6, strict_coverage=True) + assert np.isnan(res.result[2, 2]) + assert np.isfinite(res.result[N // 2, N // 2]) + + def test_default_renormalizes_corners(self): + res = symmetrize_image(HEX, 6) + assert np.isfinite(res.result[2, 2]) + + def test_nan_hole_filled_from_symmetry(self): + holed = HEX.copy() + holed[60:65, 60:65] = np.nan + res = symmetrize_image(holed, 6) + assert np.isfinite(res.result[62, 62]) + assert abs(res.result[62, 62] - HEX[62, 62]) < 0.3 + # The hole contributes nothing, so its residual is undefined there. + assert np.isnan(res.residual[62, 62]) + + +# ─── mirror (dihedral) mode ───────────────────────────────────────────────── + +class TestMirror: + def test_correct_axis_preserves_pattern(self): + lat30 = _hex_lattice(angle0_deg=30.0) # mirror-symmetric across 30° + res = symmetrize_image(lat30, 1, mirror=True, mirror_axis_deg=30.0) + assert res.n_ops == 2 + assert res.symmetry_residual_norm < 0.05 + + def test_wrong_axis_degrades(self): + lat30 = _hex_lattice(angle0_deg=30.0) + res = symmetrize_image(lat30, 1, mirror=True, mirror_axis_deg=15.0) + assert res.symmetry_residual_norm > 0.3 + + def test_d6_doubles_op_count(self): + res = symmetrize_image(HEX, 6, mirror=True, mirror_axis_deg=0.0) + assert res.n_ops == 12 + assert res.shifts.shape == (12, 2) + + +# ─── symmetrize_fft processing op ─────────────────────────────────────────── + +class TestProcessingOp: + """The op is registered end-to-end: step construction, dispatch, replay.""" + + @staticmethod + def _state(params): + from probeflow.processing.state import ProcessingState, ProcessingStep + return ProcessingState(steps=[ProcessingStep("symmetrize_fft", params)]) + + def test_step_construction_accepts_op(self): + self._state({"n_fold": 6}) + + def test_dispatch_matches_backend(self): + from probeflow.processing.state import apply_processing_state + out = apply_processing_state(HEX, self._state({"n_fold": 6})) + np.testing.assert_array_equal(out, symmetrize_filter(HEX, 6)) + + def test_default_params_are_identity(self): + from probeflow.processing.state import apply_processing_state + np.testing.assert_array_equal( + apply_processing_state(HEX, self._state({})), HEX) + + def test_replay_is_deterministic(self): + # Registration shifts are recomputed on replay; same data → same result. + from probeflow.processing.state import apply_processing_state + state = self._state({"n_fold": 6, "mirror": True, "mirror_axis_deg": 30.0}) + a = apply_processing_state(HEX, state) + b = apply_processing_state(HEX, state) + np.testing.assert_array_equal(a, b) + + def test_step_dict_round_trip(self): + from probeflow.processing.state import ProcessingStep + params = {"n_fold": 4, "mirror": True, "mirror_axis_deg": 12.5, + "register": False, "interpolation": "cubic", + "strict_coverage": True} + step = ProcessingStep.from_dict({"op": "symmetrize_fft", "params": params}) + assert step.op == "symmetrize_fft" + assert step.params == params + + def test_gui_adapter_passthrough_includes_op(self): + from probeflow.processing.gui_adapter import _FILTER_OPS_PASSTHROUGH + assert "symmetrize_fft" in _FILTER_OPS_PASSTHROUGH + + +# ─── fold_axis_from_peaks ─────────────────────────────────────────────────── + +class TestFoldAxis: + @staticmethod + def _peaks(theta0_deg, n, r=50.0): + th = np.radians(theta0_deg) + np.arange(n) * 2.0 * np.pi / n + return np.column_stack([r * np.cos(th), r * np.sin(th)]) + + def test_recovers_axis(self): + assert fold_axis_from_peaks(self._peaks(12.0, 6), 6) == pytest.approx(12.0) + + def test_result_is_modulo_fold_angle(self): + assert fold_axis_from_peaks(self._peaks(71.0, 6), 6) == pytest.approx(11.0) + + def test_robust_to_missing_peak(self): + peaks = self._peaks(12.0, 6)[:-1] + assert fold_axis_from_peaks(peaks, 6) == pytest.approx(12.0) + + def test_rejects_empty(self): + with pytest.raises(ValueError, match="non-empty"): + fold_axis_from_peaks(np.empty((0, 2)), 6) + + def test_rejects_inconsistent_angles(self): + # Two peaks exactly half a fold apart cancel in the circular mean. + peaks = self._peaks(0.0, 1)[:1] + peaks = np.vstack([peaks, self._peaks(30.0, 1)[:1]]) + with pytest.raises(ValueError, match="consistent"): + fold_axis_from_peaks(peaks, 6)