diff --git a/.github/workflows/widget-release.yml b/.github/workflows/widget-release.yml index 9b60f61f..e8acc169 100644 --- a/.github/workflows/widget-release.yml +++ b/.github/workflows/widget-release.yml @@ -129,6 +129,7 @@ jobs: wheel = wheels[0] required = { "quantem/widget/static/chooselattice.js", + "quantem/widget/static/showbragg.js", "quantem/widget/static/show1d.js", "quantem/widget/static/show2d.js", "quantem/widget/static/show3d.js", diff --git a/CHANGELOG.md b/CHANGELOG.md index 30c59833..f9b48271 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ new `rcN` heading when that rc is published to TestPyPI. ## Unreleased +- Add `ShowBragg`, an interactive front end for the `quantem.diffraction` + Bragg-disk workflow: template building, per-position correlation probing, + disk detection tuned on a sampled subset then run over the full scan, a + Bragg vector map with numbered candidates, basis selection by candidate + index or direct placement, and a per-position lattice fit handed off to + `StrainMap`. - Installation, Colab tutorial cells, and the verify snippet now install the TestPyPI wheel and import `load` from `quantem.gpu.io`. `quantem showdiffraction --demo` is documented as needing the public Fe3O4 hub folder, diff --git a/README.md b/README.md index fb729804..0017ab82 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ for backend setup, Colab instructions, and verification. | `ShowPtycho` | Interactive SSB phase and aberration review | [API](https://electronmicroscopy.github.io/quantem.widget/api/showptycho.html) | | `ShowDiffraction` | Measure diffraction spots, rings, spacing, and angles | [tutorial](https://electronmicroscopy.github.io/quantem.widget/tutorials/showdiffraction.html) · [API](https://electronmicroscopy.github.io/quantem.widget/api/showdiffraction.html) | | `ChooseLattice` | Select an origin and lattice vectors | [API](https://electronmicroscopy.github.io/quantem.widget/api/choose-lattice.html) | +| `ShowBragg` | Detect Bragg disks, pick a basis, and fit the lattice for strain | [API](https://electronmicroscopy.github.io/quantem.widget/api/showbragg.html) | | `ShowEDS` | Explore linked EDS/EELS maps and spectra | — | | `ShowFolder` | Browse, group, and select microscopy session files | [tutorial](https://electronmicroscopy.github.io/quantem.widget/tutorials/showfolder.html) · [API](https://electronmicroscopy.github.io/quantem.widget/api/showfolder.html) | diff --git a/docs/_toc.yml b/docs/_toc.yml index 5892aca7..d1d4755f 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -74,6 +74,7 @@ parts: - file: api/showeds - file: api/showdiffraction - file: api/choose-lattice + - file: api/showbragg - file: api/showfolder - file: api/html-export - file: api/io diff --git a/docs/api/index.md b/docs/api/index.md index 8e09779b..242f7574 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -37,6 +37,7 @@ readers and test agents. | [ShowPtycho](showptycho) | `quantem.widget.ShowPtycho` | ptychography state JSON, WebGPU folder export with BF-indexed `G(k)` payload | | [ShowDiffraction](showdiffraction) | `quantem.widget.showdiffraction.ShowDiffraction` | state JSON, PNG, interactive HTML | | [ChooseLattice](choose-lattice) | `quantem.widget.choose_lattice.ChooseLattice` | state JSON (`save_state=True`) | +| [ShowBragg](showbragg) | `quantem.widget.showbragg.ShowBragg` | state JSON (`save_state=True`), `BraggVectors` and `StrainMap` in Python | | [ShowFolder](showfolder) | `quantem.widget.showfolder.ShowFolder` | selection JSON for selected microscopy files and folders | All widget-level HTML exports follow the [HTML export](html-export) protocol. diff --git a/docs/api/showbragg.md b/docs/api/showbragg.md new file mode 100644 index 00000000..8948812d --- /dev/null +++ b/docs/api/showbragg.md @@ -0,0 +1,70 @@ +# ShowBragg + +Drive `quantem.diffraction.BraggVectors` Bragg-disk workflow. Probe the correlation map at a scan +position, tune disk detection on sampled subset, and from there accumulate a Bragg vector +map with fitting the lattice at every scan position. + +```python +from quantem.widget import ShowBragg + +widget = ShowBragg(dataset) # Dataset4dstem, or an existing BraggVectors +``` + +Template generation and correlation probing are updated interactively as their controls change. Operations that require processing the full scan, like Bragg-disk detection and lattice fitting, are run explicitly using buttons or Python methods. + +```python +widget.detect() # full detect_disks +widget.fit() # fit_lattice +widget.basis # (origin, g1, g2) +strain = widget.strain_map() # Construct StrainMap +``` + +Scientific results remain owned by the underlying BraggVectors workflow. The complete wrapped BraggVectors object is available through widget.bragg for functionality that is not exposed directly in the interface. + +For template_radius and max_peak_shift, a value of 0.0 means that the corresponding value should be determined automatically by quantem. + +Strain visualization is handled outside ShowBragg, the widget returns a StrainMap, which can then be visualized using the standard StrainMap.plot_strain interface. + +## Reference + +```{eval-rst} +.. autoclass:: quantem.widget.showbragg.ShowBragg + :members: + :show-inheritance: +``` + +## Display configuration + +`ShowBragg` follows the shared viewer presets described in the +[UI Guide](../developer/ui-guide): + +Each UI preset configures show_title, show_controls, and controls_collapsed. + +The controls can also be managed programmatically: + +* widget.collapse_controls() +* widget.expand_controls() +* widget.toggle_controls() + +When the controls are collapsed, parameter columns and execution buttons are hidden while the primary visualization remains visible. + +## Interactive controls + +| Control | Trait | Behavior | +|---|---|---| +| Template source selector | `template_source` | Rebuilds the template from a synthetic disk, the data mean, or a supplied probe | +| Template radius / edge | `template_radius`, `template_edge` | Re-renders `template_png` and the correlation map | +| Subtract mean toggle | `template_subtract_mean` | Rebuilds the template as a zero-sum band-pass kernel | +| Scan row / col fields | `probe_position` | Re-renders the diffraction pattern and its correlation map | +| Six detection fields | `min_abs_intensity`, `min_spacing`, `edge_boundary`, `subpixel`, `upsample_factor`, `max_num_peaks` | Stored for the next preview or full run | +| Preview on grid button | `preview_peaks`, `detection_state` | Detects on `preview_grid` x `preview_grid` sampled positions and marks them on the probe panel | +| Run full detection button | `detection_state`, `bvm_png`, `candidates` | Detects at every scan position, then accumulates the Bragg vector map | +| Recompute map button | `bvm_sampling`, `bvm_png` | Re-accumulates the Bragg vector map at the given sampling | +| Candidate fields | `num_candidates`, `candidate_min_spacing`, `candidate_min_abs_intensity` | Re-derives the numbered candidate set live | +| Click a numbered candidate | `origin_index`, `g1_index`, `g2_index` | Assigns that candidate to the active basis role | +| Drag a basis marker | `origin_rc`, `g1_rc`, `g2_rc` | Places the vector freely and clears that role's candidate index | +| Reset to automatic button | all six basis traits | Lets quantem pick the basis again | +| Run fit button | `fit_state`, `mask_weight_png`, `fit_error_png` | Fits the lattice at every position and shows the two diagnostics | +| Controls / Hide button | `controls_collapsed` | Hides the parameter columns and run buttons, keeping the images | +| Pan (drag) / zoom (wheel) | view transform | Any image panel translates / zooms about the cursor | +| Double-click | view transform | Resets that panel's zoom and pan | diff --git a/docs/developer/ui-guide.md b/docs/developer/ui-guide.md index 18dc0239..4b1329ba 100644 --- a/docs/developer/ui-guide.md +++ b/docs/developer/ui-guide.md @@ -28,12 +28,12 @@ Show3D( | Name | Widgets | Meaning | |---|---|---| -| `show_title` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction | Show or hide the top title row. | -| `show_controls` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction | Expose controls at all. Set `False` for a permanently clean display with no GUI recovery button. | -| `controls_collapsed` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction | Start with controls hidden. Show4DSTEM exposes this as programmatic/state control only; other widgets may also show a top `Controls` button. | -| `collapse_controls()` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction | Programmatically collapse controls. | -| `expand_controls()` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction | Programmatically expand controls. | -| `toggle_controls()` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction | Programmatically toggle collapsed controls. | +| `show_title` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction, ShowBragg | Show or hide the top title row. | +| `show_controls` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction, ShowBragg | Expose controls at all. Set `False` for a permanently clean display with no GUI recovery button. | +| `controls_collapsed` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction, ShowBragg | Start with controls hidden. Show4DSTEM exposes this as programmatic/state control only; other widgets may also show a top `Controls` button. | +| `collapse_controls()` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction, ShowBragg | Programmatically collapse controls. | +| `expand_controls()` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction, ShowBragg | Programmatically expand controls. | +| `toggle_controls()` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowEDS, ShowDiffraction, ShowBragg | Programmatically toggle collapsed controls. | | `show_stats` | Show1D, Show2D, Show3D, Show3DSlices, Show4DSTEM, ShowDiffraction | Show or hide mean/min/max/std readouts where the widget renders or computes them. | | `show_panel_titles` | Show2D, Show3D | Show labels on individual image/volume panels. | | `panel_title_font_size` | Show2D, Show3D | Font size for per-panel labels. | diff --git a/js/chooselattice/index.tsx b/js/chooselattice/index.tsx index 88cdcb1a..6773617c 100644 --- a/js/chooselattice/index.tsx +++ b/js/chooselattice/index.tsx @@ -14,11 +14,19 @@ import Typography from "@mui/material/Typography"; import Stack from "@mui/material/Stack"; import Button from "@mui/material/Button"; import { useTheme } from "../theme"; -import { extractBytes, preserveRestoredWidgetModelsOnSave } from "../format"; +import { preserveRestoredWidgetModelsOnSave } from "../format"; import { useHideStaticFallback } from "../staticFallback"; +import { + canvasPoint, + clamp, + drawImage, + imageToScreen, + screenToImage, + usePngBitmap, + zoomAt, + type ImageViewport, +} from "../imageView"; -const MIN_ZOOM = 0.5; -const MAX_ZOOM = 20; const CANVAS_SIZE = 512; const CANVAS_BORDER_PX = 1; const HIT_PX = 10; @@ -37,10 +45,6 @@ const compactButton = { type Point = [number, number]; // [row, col] in original image pixels type DragMode = "none" | "pan" | "point"; -function clamp(value: number, lo: number, hi: number): number { - return Math.min(hi, Math.max(lo, value)); -} - function ChooseLattice() { const model = useModel(); const rootRef = React.useRef(null); @@ -56,38 +60,19 @@ function ChooseLattice() { const [pointLabels] = useModelState("point_labels"); const [points, setPoints] = useModelState("points"); - // Decode the PNG payload once per change into a drawable bitmap. - const [image, setImage] = React.useState(null); - React.useEffect(() => { - const bytes = extractBytes(frameBytes); - if (bytes.length === 0) { - setImage(null); - return; - } - let cancelled = false; - const blob = new Blob([bytes as unknown as BlobPart], { type: "image/png" }); - if (typeof createImageBitmap === "function") { - createImageBitmap(blob).then((bmp) => { if (!cancelled) setImage(bmp); }); - } else { - const url = URL.createObjectURL(blob); - const img = new Image(); - img.onload = () => { if (!cancelled) setImage(img); URL.revokeObjectURL(url); }; - img.src = url; - } - return () => { cancelled = true; }; - }, [frameBytes]); + const image = usePngBitmap(frameBytes); // View state: zoom + pan (CSS px, canvas-centered). const [zoom, setZoom] = React.useState(1); const [panX, setPanX] = React.useState(0); const [panY, setPanY] = React.useState(0); - // displayScale maps original image pixels -> CSS px at zoom=1. - const displayScale = height > 0 && width > 0 - ? CANVAS_SIZE / Math.max(height, width) - : 1; const canvasW = CANVAS_SIZE; const canvasH = CANVAS_SIZE; + const viewport: ImageViewport = React.useMemo( + () => ({ height, width, canvas: CANVAS_SIZE, zoom, panX, panY }), + [height, width, zoom, panX, panY], + ); const canvasRef = React.useRef(null); const uiRef = React.useRef(null); @@ -116,13 +101,7 @@ function ChooseLattice() { ctx.fillStyle = themeColors.bg; ctx.fillRect(0, 0, canvasW, canvasH); if (!image || !width || !height) return; - const cx = canvasW / 2; - const cy = canvasH / 2; - const drawW = width * displayScale * zoom; - const drawH = height * displayScale * zoom; - const x = cx - drawW / 2 + panX; - const y = cy - drawH / 2 + panY; - ctx.drawImage(image, x, y, drawW, drawH); + drawImage(ctx, image, viewport); // Confirm the decoded bitmap on the next compositing frame. Static docs // can mount while Chrome is still promoting the canvas layer; without a // second paint that one-shot draw can remain a white presentation frame. @@ -132,42 +111,33 @@ function ChooseLattice() { if (!current || !currentCtx) return; currentCtx.fillStyle = themeColors.bg; currentCtx.fillRect(0, 0, canvasW, canvasH); - currentCtx.drawImage(image, x, y, drawW, drawH); + drawImage(currentCtx, image, viewport); }); return () => window.cancelAnimationFrame(confirmFrame); - }, [image, width, height, displayScale, zoom, panX, panY, canvasW, canvasH, themeColors.bg]); + }, [image, width, height, viewport, canvasW, canvasH, themeColors.bg]); // Convert a mouse event to original-image (row, col) coordinates. const screenToImg = React.useCallback((e: { clientX: number; clientY: number }): Point => { const canvas = canvasRef.current; if (!canvas) return [0, 0]; - const rect = canvas.getBoundingClientRect(); - const mouseCanvasX = (e.clientX - rect.left) * (canvas.width / rect.width); - const mouseCanvasY = (e.clientY - rect.top) * (canvas.height / rect.height); - const cx = canvasW / 2; - const cy = canvasH / 2; - const col = (mouseCanvasX - cx - panX) / (displayScale * zoom) + width / 2; - const row = (mouseCanvasY - cy - panY) / (displayScale * zoom) + height / 2; - return [row, col]; - }, [canvasW, canvasH, panX, panY, displayScale, zoom, width, height]); - - const imgToScreen = React.useCallback((row: number, col: number): [number, number] => { - const cx = canvasW / 2; - const cy = canvasH / 2; - const x = cx + (col - width / 2) * displayScale * zoom + panX; - const y = cy + (row - height / 2) * displayScale * zoom + panY; - return [x, y]; - }, [canvasW, canvasH, panX, panY, displayScale, zoom, width, height]); + const [x, y] = canvasPoint(canvas, e); + return screenToImage(viewport, x, y); + }, [viewport]); + + const imgToScreen = React.useCallback( + (row: number, col: number): [number, number] => imageToScreen(viewport, row, col), + [viewport], + ); const hitTestPoint = React.useCallback((row: number, col: number): number => { - const hitArea = HIT_PX / (displayScale * zoom); + const hitArea = HIT_PX / ((canvasW / Math.max(height, width)) * zoom); const list = points || []; for (let i = list.length - 1; i >= 0; i--) { const [pr, pc] = list[i]; if (Math.hypot(row - pr, col - pc) <= hitArea) return i; } return -1; - }, [points, displayScale, zoom]); + }, [points, canvasW, height, width, zoom]); // Wheel: cursor-anchored zoom. Page-scroll prevention is handled by a // native non-passive listener below (React's synthetic onWheel is passive, @@ -175,18 +145,11 @@ function ChooseLattice() { const handleWheel = (e: React.WheelEvent) => { const canvas = canvasRef.current; if (!canvas) return; - const rect = canvas.getBoundingClientRect(); - const mouseCanvasX = (e.clientX - rect.left) * (canvas.width / rect.width); - const mouseCanvasY = (e.clientY - rect.top) * (canvas.height / rect.height); - const cx = canvasW / 2; - const cy = canvasH / 2; - const mouseImageX = (mouseCanvasX - cx - panX) / zoom + cx; - const mouseImageY = (mouseCanvasY - cy - panY) / zoom + cy; - const zoomFactor = e.deltaY > 0 ? 0.9 : 1.1; - const newZoom = clamp(zoom * zoomFactor, MIN_ZOOM, MAX_ZOOM); - setPanX(mouseCanvasX - (mouseImageX - cx) * newZoom - cx); - setPanY(mouseCanvasY - (mouseImageY - cy) * newZoom - cy); - setZoom(newZoom); + const [x, y] = canvasPoint(canvas, e); + const next = zoomAt(viewport, x, y, e.deltaY); + setPanX(next.panX); + setPanY(next.panY); + setZoom(next.zoom); }; const resetView = React.useCallback(() => { diff --git a/js/imageView.test.ts b/js/imageView.test.ts new file mode 100644 index 00000000..47115dee --- /dev/null +++ b/js/imageView.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it } from "vitest"; +import { imageToScreen, screenToImage, zoomAt, type ImageViewport } from "./imageView"; + +const view: ImageViewport = { height: 48, width: 64, canvas: 512, zoom: 1, panX: 0, panY: 0 }; + +describe("imageView", () => { + it("round-trips image and screen coordinates under pan and zoom", () => { + const panned: ImageViewport = { ...view, zoom: 3.2, panX: -40, panY: 17 }; + const [x, y] = imageToScreen(panned, 12.5, 30.25); + const [row, col] = screenToImage(panned, x, y); + expect(row).toBeCloseTo(12.5); + expect(col).toBeCloseTo(30.25); + }); + + it("keeps the point under the cursor fixed while zooming", () => { + const before = screenToImage(view, 300, 200); + const zoomed: ImageViewport = { ...view, ...zoomAt(view, 300, 200, -1) }; + const after = screenToImage(zoomed, 300, 200); + expect(after[0]).toBeCloseTo(before[0]); + expect(after[1]).toBeCloseTo(before[1]); + }); +}); diff --git a/js/imageView.ts b/js/imageView.ts new file mode 100644 index 00000000..3515d62e --- /dev/null +++ b/js/imageView.ts @@ -0,0 +1,135 @@ +/** Shared pan/zoom transform for the single-canvas image widgets. */ + +import { useEffect, useState } from "react"; +import { extractBytes } from "./format"; + +export const MIN_ZOOM = 0.5; +export const MAX_ZOOM = 20; + +export interface ImageViewport { + /** Original image rows and columns. */ + height: number; + width: number; + /** Square canvas edge, in CSS px. */ + canvas: number; + zoom: number; + panX: number; + panY: number; +} + +export type ViewTransform = Pick; + +export const IDENTITY_VIEW: ViewTransform = { zoom: 1, panX: 0, panY: 0 }; + +export function clamp(value: number, lo: number, hi: number): number { + return Math.min(hi, Math.max(lo, value)); +} + +/** Image pixels per CSS px at zoom 1, fitting the longest side to the canvas. */ +export function displayScale(v: ImageViewport): number { + const longest = Math.max(v.height, v.width); + return longest > 0 ? v.canvas / longest : 1; +} + +/** Canvas coordinates of a mouse event, corrected for CSS scaling. */ +export function canvasPoint( + canvas: HTMLCanvasElement, + e: { clientX: number; clientY: number }, +): [number, number] { + const rect = canvas.getBoundingClientRect(); + return [ + (e.clientX - rect.left) * (canvas.width / rect.width), + (e.clientY - rect.top) * (canvas.height / rect.height), + ]; +} + +export function screenToImage( + v: ImageViewport, + canvasX: number, + canvasY: number, +): [number, number] { + const scale = displayScale(v) * v.zoom; + const center = v.canvas / 2; + const col = (canvasX - center - v.panX) / scale + v.width / 2; + const row = (canvasY - center - v.panY) / scale + v.height / 2; + return [row, col]; +} + +export function imageToScreen( + v: ImageViewport, + row: number, + col: number, +): [number, number] { + const scale = displayScale(v) * v.zoom; + const center = v.canvas / 2; + return [ + center + (col - v.width / 2) * scale + v.panX, + center + (row - v.height / 2) * scale + v.panY, + ]; +} + +/** Wheel zoom that keeps the image point under the cursor fixed. */ +export function zoomAt( + v: ImageViewport, + canvasX: number, + canvasY: number, + deltaY: number, +): ViewTransform { + const center = v.canvas / 2; + const anchorX = (canvasX - center - v.panX) / v.zoom + center; + const anchorY = (canvasY - center - v.panY) / v.zoom + center; + const zoom = clamp(v.zoom * (deltaY > 0 ? 0.9 : 1.1), MIN_ZOOM, MAX_ZOOM); + return { + zoom, + panX: canvasX - (anchorX - center) * zoom - center, + panY: canvasY - (anchorY - center) * zoom - center, + }; +} + +/** Paint the image into the canvas under the current view. */ +export function drawImage( + ctx: CanvasRenderingContext2D, + image: CanvasImageSource, + v: ImageViewport, +): void { + const scale = displayScale(v) * v.zoom; + const center = v.canvas / 2; + const drawW = v.width * scale; + const drawH = v.height * scale; + ctx.drawImage(image, center - drawW / 2 + v.panX, center - drawH / 2 + v.panY, drawW, drawH); +} + +/** Decode PNG bytes from a synced trait into a drawable bitmap. */ +export function usePngBitmap( + bytes: DataView | Uint8Array | null | undefined, +): ImageBitmap | HTMLImageElement | null { + const [image, setImage] = useState(null); + + useEffect(() => { + const raw = bytes ? extractBytes(bytes) : new Uint8Array(0); + if (raw.length === 0) { + setImage(null); + return; + } + let cancelled = false; + const blob = new Blob([raw as unknown as BlobPart], { type: "image/png" }); + if (typeof createImageBitmap === "function") { + createImageBitmap(blob).then((bitmap) => { + if (!cancelled) setImage(bitmap); + }); + } else { + const url = URL.createObjectURL(blob); + const element = new Image(); + element.onload = () => { + if (!cancelled) setImage(element); + URL.revokeObjectURL(url); + }; + element.src = url; + } + return () => { + cancelled = true; + }; + }, [bytes]); + + return image; +} diff --git a/js/showbragg/index.tsx b/js/showbragg/index.tsx new file mode 100644 index 00000000..f5f6a035 --- /dev/null +++ b/js/showbragg/index.tsx @@ -0,0 +1,905 @@ +/** ShowBragg - panel stack driving the quantem BraggVectors workflow. */ + +import * as React from "react"; +import { createRender, useModel, useModelState } from "@anywidget/react"; +import Box from "@mui/material/Box"; +import Button from "@mui/material/Button"; +import MenuItem from "@mui/material/MenuItem"; +import Select from "@mui/material/Select"; +import Stack from "@mui/material/Stack"; +import Switch from "@mui/material/Switch"; +import FormControlLabel from "@mui/material/FormControlLabel"; +import TextField from "@mui/material/TextField"; +import Typography from "@mui/material/Typography"; +import { useTheme, type ThemeColors } from "../theme"; +import { preserveRestoredWidgetModelsOnSave } from "../format"; +import { useHideStaticFallback } from "../staticFallback"; +import { + canvasPoint, + clamp, + drawImage, + imageToScreen, + screenToImage, + usePngBitmap, + zoomAt, + type ImageViewport, + type ViewTransform, +} from "../imageView"; + +const CANVAS_SIZE = 340; +const HIT_PX = 10; +const CLICK_MOVE_THRESHOLD_PX = 4; +const BASIS_COLORS = ["#ff4d4f", "#40a9ff", "#73d13d"]; +const LABEL_CLEARANCE_PX = 26; +const BASIS_LABELS = ["origin", "g1", "g2"] as const; +const SPACING = { XS: 4, SM: 8, MD: 12, LG: 16 } as const; +const CMAPS = ["inferno", "magma", "viridis", "plasma", "gray", "turbo"]; + +const compactButton = { + fontSize: 10, + py: 0.25, + px: 1, + minWidth: 0, + textTransform: "none" as const, +}; + +type Role = (typeof BASIS_LABELS)[number]; +type Marker = { row: number; col: number; color: string; label: string }; +type PreviewPeaks = { positions: number[][]; counts: number[]; peaks: number[][] }; + +function median(values: number[]): number { + if (!values.length) return 0; + const sorted = [...values].sort((a, b) => a - b); + return sorted[Math.floor(sorted.length / 2)]; +} + +function Panel({ + title, + colors, + blocked, + children, +}: { + title: string; + colors: ThemeColors; + blocked?: string; + children: React.ReactNode; +}) { + const [open, setOpen] = React.useState(true); + + return ( + + setOpen(!open)} + sx={{ px: `${SPACING.MD}px`, py: `${SPACING.SM}px`, cursor: "pointer", bgcolor: colors.bgAlt }} + > + {title} + {open ? "hide" : "show"} + + + {open && ( + + {blocked ? ( + {blocked} + ) : ( + children + )} + + )} + + ); +} + +function NumberField({ + label, + value, + onChange, + step = 1, + disabled, + colors, + hint, +}: { + label: string; + value: number; + onChange: (value: number) => void; + step?: number; + disabled?: boolean; + colors: ThemeColors; + hint?: string; +}) { + return ( + { + const parsed = Number(e.target.value); + if (Number.isFinite(parsed)) onChange(parsed); + }} + slotProps={{ htmlInput: { step }, inputLabel: { shrink: true } }} + sx={{ + width: 120, + "& .MuiInputBase-input": { fontSize: 11, py: 0.5, color: colors.text }, + "& .MuiInputLabel-root": { fontSize: 11, color: colors.textMuted }, + }} + /> + ); +} + +function ImageCanvas({ + bytes, + shape, + colors, + overlay, + markers, + onMoveMarker, + onCommitMarker, + onPick, + caption, +}: { + bytes: DataView | null; + shape: number[]; + colors: ThemeColors; + overlay?: (ctx: CanvasRenderingContext2D, viewport: ImageViewport) => void; + markers?: Marker[]; + onMoveMarker?: (index: number, row: number, col: number) => void; + onCommitMarker?: (index: number, row: number, col: number) => void; + onPick?: (row: number, col: number) => void; + caption?: string; +}) { + const image = usePngBitmap(bytes); + const [view, setView] = React.useState({ zoom: 1, panX: 0, panY: 0 }); + const baseRef = React.useRef(null); + const overlayRef = React.useRef(null); + const containerRef = React.useRef(null); + + const [height, width] = shape.length === 2 ? shape : [1, 1]; + const viewport: ImageViewport = React.useMemo( + () => ({ height, width, canvas: CANVAS_SIZE, ...view }), + [height, width, view], + ); + + React.useEffect(() => { + const el = containerRef.current; + if (!el) return; + const prevent = (e: WheelEvent) => e.preventDefault(); + el.addEventListener("wheel", prevent, { passive: false }); + return () => el.removeEventListener("wheel", prevent); + }, []); + + React.useLayoutEffect(() => { + const canvas = baseRef.current; + const ctx = canvas?.getContext("2d"); + if (!canvas || !ctx) return; + canvas.width = CANVAS_SIZE; + canvas.height = CANVAS_SIZE; + ctx.imageSmoothingEnabled = false; + ctx.fillStyle = colors.bg; + ctx.fillRect(0, 0, CANVAS_SIZE, CANVAS_SIZE); + if (image) drawImage(ctx, image, viewport); + }, [image, viewport, colors.bg]); + + React.useLayoutEffect(() => { + const canvas = overlayRef.current; + const ctx = canvas?.getContext("2d"); + if (!canvas || !ctx) return; + canvas.width = CANVAS_SIZE; + canvas.height = CANVAS_SIZE; + ctx.clearRect(0, 0, CANVAS_SIZE, CANVAS_SIZE); + overlay?.(ctx, viewport); + + (markers ?? []).forEach(({ row, col, color, label }) => { + const [x, y] = imageToScreen(viewport, row, col); + ctx.beginPath(); + ctx.arc(x, y, 5, 0, 2 * Math.PI); + ctx.fillStyle = color; + ctx.fill(); + ctx.strokeStyle = "#000"; + ctx.lineWidth = 1; + ctx.stroke(); + ctx.font = "bold 11px -apple-system, sans-serif"; + ctx.fillStyle = "#fff"; + ctx.strokeStyle = "rgba(0,0,0,0.85)"; + ctx.lineWidth = 3; + ctx.strokeText(label, x + 8, y - 6); + ctx.fillText(label, x + 8, y - 6); + }); + }, [overlay, markers, viewport]); + + const dragRef = React.useRef<{ + moved: boolean; + startX: number; + startY: number; + panX: number; + panY: number; + marker: number; + } | null>(null); + + const toImage = (e: React.MouseEvent): [number, number] => { + const canvas = baseRef.current; + if (!canvas) return [0, 0]; + const [x, y] = canvasPoint(canvas, e); + return screenToImage(viewport, x, y); + }; + + const hitMarker = (row: number, col: number): number => { + if (!onMoveMarker) return -1; + const scale = (CANVAS_SIZE / Math.max(height, width)) * view.zoom; + const radius = HIT_PX / scale; + const list = markers ?? []; + for (let i = list.length - 1; i >= 0; i--) { + if (Math.hypot(row - list[i].row, col - list[i].col) <= radius) return i; + } + return -1; + }; + + const handleMouseDown = (e: React.MouseEvent) => { + const [row, col] = toImage(e); + dragRef.current = { + moved: false, + startX: e.clientX, + startY: e.clientY, + panX: view.panX, + panY: view.panY, + marker: hitMarker(row, col), + }; + }; + + const handleMouseMove = (e: React.MouseEvent) => { + const drag = dragRef.current; + if (!drag) return; + if (Math.hypot(e.clientX - drag.startX, e.clientY - drag.startY) > CLICK_MOVE_THRESHOLD_PX) { + drag.moved = true; + } + if (drag.marker >= 0) { + const [row, col] = toImage(e); + onMoveMarker?.( + drag.marker, + clamp(row, 0, height - 1), + clamp(col, 0, width - 1), + ); + return; + } + if (drag.moved) { + setView((v) => ({ + ...v, + panX: drag.panX + (e.clientX - drag.startX), + panY: drag.panY + (e.clientY - drag.startY), + })); + } + }; + + const handleMouseUp = (e: React.MouseEvent) => { + const drag = dragRef.current; + dragRef.current = null; + if (!drag) return; + if (drag.marker >= 0) { + const [row, col] = toImage(e); + onCommitMarker?.( + drag.marker, + clamp(row, 0, height - 1), + clamp(col, 0, width - 1), + ); + return; + } + if (drag.moved) return; + const [row, col] = toImage(e); + onPick?.(row, col); + }; + + return ( + + + + + { + const canvas = baseRef.current; + if (!canvas) return; + const [x, y] = canvasPoint(canvas, e); + setView(zoomAt(viewport, x, y, e.deltaY)); + }} + onMouseDown={handleMouseDown} + onMouseMove={handleMouseMove} + onMouseUp={handleMouseUp} + onMouseLeave={() => { dragRef.current = null; }} + onDoubleClick={() => setView({ zoom: 1, panX: 0, panY: 0 })} + /> + + + {caption && ( + + {caption} + + )} + + ); +} + +function ShowBragg() { + const model = useModel(); + const rootRef = React.useRef(null); + const { colors } = useTheme(); + + React.useEffect(() => preserveRestoredWidgetModelsOnSave(model), [model]); + useHideStaticFallback(model, rootRef); + + const [title] = useModelState("title"); + const [showTitle] = useModelState("show_title"); + const [showControls] = useModelState("show_controls"); + const [controlsCollapsed, setControlsCollapsed] = useModelState("controls_collapsed"); + const controlsVisible = showControls && !controlsCollapsed; + const [scanShape] = useModelState("scan_shape"); + const [qShape] = useModelState("q_shape"); + const [status] = useModelState("status"); + const [cmap, setCmap] = useModelState("cmap"); + const [logScale, setLogScale] = useModelState("log_scale"); + + const [templateSource, setTemplateSource] = useModelState("template_source"); + const [templateRadius, setTemplateRadius] = useModelState("template_radius"); + const [templateEdge, setTemplateEdge] = useModelState("template_edge"); + const [templateSubtractMean, setTemplateSubtractMean] = useModelState("template_subtract_mean"); + const [templatePng] = useModelState("template_png"); + const [templateShape] = useModelState("template_shape"); + const [hasProbe] = useModelState("has_probe"); + + const [probePosition, setProbePosition] = useModelState("probe_position"); + const [probeDiffractionPng] = useModelState("probe_diffraction_png"); + const [probeCorrelationPng] = useModelState("probe_correlation_png"); + + const [minAbsIntensity, setMinAbsIntensity] = useModelState("min_abs_intensity"); + const [minSpacing, setMinSpacing] = useModelState("min_spacing"); + const [edgeBoundary, setEdgeBoundary] = useModelState("edge_boundary"); + const [subpixel, setSubpixel] = useModelState("subpixel"); + const [upsampleFactor, setUpsampleFactor] = useModelState("upsample_factor"); + const [maxNumPeaks, setMaxNumPeaks] = useModelState("max_num_peaks"); + const [previewGrid, setPreviewGrid] = useModelState("preview_grid"); + const [previewPeaksJson] = useModelState("preview_peaks"); + const [detectionState] = useModelState("detection_state"); + + const [bvmSampling, setBvmSampling] = useModelState("bvm_sampling"); + const [bvmPng] = useModelState("bvm_png"); + const [numCandidates, setNumCandidates] = useModelState("num_candidates"); + const [candidateMinSpacing, setCandidateMinSpacing] = useModelState("candidate_min_spacing"); + const [candidateMinAbsIntensity, setCandidateMinAbsIntensity] = useModelState("candidate_min_abs_intensity"); + const [candidates] = useModelState("candidates"); + const [originIndex, setOriginIndex] = useModelState("origin_index"); + const [g1Index, setG1Index] = useModelState("g1_index"); + const [g2Index, setG2Index] = useModelState("g2_index"); + const [originRc, setOriginRc] = useModelState("origin_rc"); + const [g1Rc, setG1Rc] = useModelState("g1_rc"); + const [g2Rc, setG2Rc] = useModelState("g2_rc"); + + const [minNumPeaks, setMinNumPeaks] = useModelState("min_num_peaks"); + const [maxPeakShift, setMaxPeakShift] = useModelState("max_peak_shift"); + const [fitState] = useModelState("fit_state"); + const [maskWeightPng] = useModelState("mask_weight_png"); + const [fitErrorPng] = useModelState("fit_error_png"); + + const send = React.useCallback( + (content: Record) => model.send(content), + [model], + ); + + const detecting = detectionState === "running" || detectionState === "preview"; + const fitting = fitState === "running"; + const busy = detecting || fitting; + const hasPeaks = candidates.length > 0 || bvmPng?.byteLength > 0; + + const preview: PreviewPeaks | null = React.useMemo(() => { + if (!previewPeaksJson) return null; + try { + return JSON.parse(previewPeaksJson) as PreviewPeaks; + } catch { + return null; + } + }, [previewPeaksJson]); + + const drawPreviewPeaks = React.useCallback( + (ctx: CanvasRenderingContext2D, viewport: ImageViewport) => { + if (!preview) return; + ctx.strokeStyle = "#40a9ff"; + ctx.lineWidth = 1.5; + preview.peaks.forEach(([row, col]) => { + const [x, y] = imageToScreen(viewport, row, col); + ctx.beginPath(); + ctx.moveTo(x - 5, y); + ctx.lineTo(x + 5, y); + ctx.moveTo(x, y - 5); + ctx.lineTo(x, y + 5); + ctx.stroke(); + }); + }, + [preview], + ); + + const drawCandidates = React.useCallback( + (ctx: CanvasRenderingContext2D, viewport: ImageViewport) => { + const placed: Array<[number, number]> = []; + const byBrightness = candidates + .map((c, i) => ({ row: c[0], col: c[1], intensity: c[2] ?? 0, index: i })) + .sort((a, b) => b.intensity - a.intensity); + + ctx.font = "10px -apple-system, sans-serif"; + for (const { row, col, index } of byBrightness) { + const [x, y] = imageToScreen(viewport, row, col); + ctx.beginPath(); + ctx.arc(x, y, 3, 0, 2 * Math.PI); + ctx.strokeStyle = "rgba(255,255,255,0.8)"; + ctx.lineWidth = 1; + ctx.stroke(); + + const clear = placed.every( + ([px, py]) => Math.hypot(x - px, y - py) >= LABEL_CLEARANCE_PX, + ); + if (!clear) continue; + placed.push([x, y]); + + ctx.fillStyle = "#fff"; + ctx.strokeStyle = "rgba(0,0,0,0.85)"; + ctx.lineWidth = 2.5; + ctx.strokeText(String(index), x + 4, y - 4); + ctx.fillText(String(index), x + 4, y - 4); + } + }, + [candidates], + ); + + const [dragPreview, setDragPreview] = React.useState<[number, number, number] | null>(null); + + const basisMarkers: Marker[] = React.useMemo(() => { + if (originRc.length !== 2) return []; + const tips = [originRc, g1Rc, g2Rc]; + return tips.flatMap((rc, i) => { + if (rc.length !== 2) return []; + if (dragPreview && dragPreview[0] === i) { + return [{ row: dragPreview[1], col: dragPreview[2], color: BASIS_COLORS[i], label: BASIS_LABELS[i] }]; + } + const row = i === 0 ? rc[0] : originRc[0] + rc[0]; + const col = i === 0 ? rc[1] : originRc[1] + rc[1]; + return [{ row, col, color: BASIS_COLORS[i], label: BASIS_LABELS[i] }]; + }); + }, [originRc, g1Rc, g2Rc, dragPreview]); + + const setRoleVector = (index: number, row: number, col: number) => { + setDragPreview(null); + if (index === 0) { + setOriginIndex(-1); + setOriginRc([row, col]); + } else if (index === 1) { + setG1Index(-1); + setG1Rc([row - originRc[0], col - originRc[1]]); + } else { + setG2Index(-1); + setG2Rc([row - originRc[0], col - originRc[1]]); + } + }; + + const [activeRole, setActiveRole] = React.useState("origin"); + + const pickCandidate = (row: number, col: number) => { + if (!candidates.length) return; + let nearest = 0; + let best = Infinity; + candidates.forEach(([cr, cc], i) => { + const d = Math.hypot(row - cr, col - cc); + if (d < best) { + best = d; + nearest = i; + } + }); + if (activeRole === "origin") setOriginIndex(nearest); + else if (activeRole === "g1") setG1Index(nearest); + else setG2Index(nearest); + }; + + const vectorText = (rc: number[]) => + rc.length === 2 ? `(${rc[0].toFixed(2)}, ${rc[1].toFixed(2)})` : "not set"; + + return ( + + {(showTitle || showControls) && ( + + {showTitle ? ( + {title || "Show Bragg"} + ) : ( + + )} + + {showTitle && ( + + scan {scanShape.join(" x ")} · detector {qShape.join(" x ")} + + )} + {showControls && ( + + )} + + + )} + + {status && ( + {status} + )} + + {controlsVisible && ( + + + DISPLAY + + + setLogScale(e.target.checked)} + /> + } + label="log scale" + slotProps={{ typography: { sx: { fontSize: 11, color: colors.text } } }} + /> + + scroll to zoom · drag to pan · double-click to reset + + + )} + + + + {controlsVisible && ( + + + + + + setTemplateSubtractMean(e.target.checked)} + /> + } + label="Subtract Mean" + slotProps={{ typography: { sx: { fontSize: 11, color: colors.text } } }} + sx={{ alignSelf: "flex-start", ml: 0 }} + /> + + )} + + 0 ? templateShape : qShape} + colors={colors} + caption="correlation template, cropped to the disk" + /> + + + + + + {controlsVisible && ( + + setProbePosition([v, probePosition[1]])} + colors={colors} + /> + setProbePosition([probePosition[0], v])} + colors={colors} + /> + + )} + + + + + + + + {controlsVisible && ( + + + + + + + + + + )} + + {controlsVisible && ( + + + + + )} + + {preview && ( + + {preview.peaks.length} peaks at the probed position · across {preview.counts.length} sampled + positions min {Math.min(...preview.counts)}, median {median(preview.counts)}, max{" "} + {Math.max(...preview.counts)}. Marks are drawn on the probe panel. + + )} + + + + + {controlsVisible && ( + + + + + )} + + 0 ? "" : " (auto)"}, brightest numbered; zoom in for more`} + /> + + + + + + {controlsVisible && ( + + + + + + + {BASIS_LABELS.map((role, i) => ( + + ))} + + + + + + + origin {vectorText(originRc)}{originIndex >= 0 ? ` · candidate ${originIndex}` : ""} + + + g1 {vectorText(g1Rc)}{g1Index >= 0 ? ` · candidate ${g1Index}` : ""} + + + g2 {vectorText(g2Rc)}{g2Index >= 0 ? ` · candidate ${g2Index}` : ""} + + + + )} + + setDragPreview([i, row, col])} + onCommitMarker={setRoleVector} + onPick={pickCandidate} + caption={`click a numbered candidate to set ${activeRole}, or drag a marker`} + /> + + + + + {controlsVisible && ( + + + + + + )} + + {fitState === "done" && ( + + + + + )} + + + ); +} + +export const render = createRender(ShowBragg); diff --git a/scripts/build.mjs b/scripts/build.mjs index dc3327aa..cbe962f8 100644 --- a/scripts/build.mjs +++ b/scripts/build.mjs @@ -21,6 +21,7 @@ const widgets = [ { name: "showeds" }, { name: "showptycho" }, { name: "chooselattice" }, + { name: "showbragg" }, ]; rmSync("src/quantem/widget/static", { recursive: true, force: true }); diff --git a/scripts/widget_release_check.sh b/scripts/widget_release_check.sh index 10f3f6be..37df9c85 100755 --- a/scripts/widget_release_check.sh +++ b/scripts/widget_release_check.sh @@ -72,6 +72,7 @@ if len(wheels) != 1: wheel = wheels[0] required = { "quantem/widget/static/chooselattice.js", + "quantem/widget/static/showbragg.js", "quantem/widget/static/show1d.js", "quantem/widget/static/show2d.js", "quantem/widget/static/show3d.js", diff --git a/src/quantem/widget/__init__.py b/src/quantem/widget/__init__.py index 075e7ccd..5a09b8b1 100644 --- a/src/quantem/widget/__init__.py +++ b/src/quantem/widget/__init__.py @@ -19,6 +19,7 @@ _LAZY_EXPORTS: dict[str, tuple[str, str | None]] = { "ChooseLattice": ("quantem.widget.choose_lattice", "ChooseLattice"), "Mask2D": ("quantem.widget.mask2d", "Mask2D"), + "ShowBragg": ("quantem.widget.showbragg", "ShowBragg"), "Show1D": ("quantem.widget.show1d", "Show1D"), "Show2D": ("quantem.widget.show2d", "Show2D"), "Show3D": ("quantem.widget.show3d", "Show3D"), diff --git a/src/quantem/widget/choose_lattice.py b/src/quantem/widget/choose_lattice.py index ec4a44a3..378d6ed1 100644 --- a/src/quantem/widget/choose_lattice.py +++ b/src/quantem/widget/choose_lattice.py @@ -12,9 +12,10 @@ from typing import Any, Sequence import anywidget -import matplotlib import numpy as np import traitlets +from quantem.widget.render import frame_to_rgb, rgb_to_png_bytes +from quantem.widget.utils.traits import reject_unknown_kwargs from quantem.widget.utils.array import to_numpy from quantem.widget.utils.static_fallback import StaticFallbackMixin @@ -36,45 +37,6 @@ def _core_image_dataset_types() -> tuple[type[Any], ...]: return _CORE_IMAGE_DATASET_TYPES -def _reject_unknown_kwargs(cls, kwargs: dict) -> None: - """Raise TypeError for any kwarg that isn't a declared trait (catches typos).""" - traits = set(cls.class_trait_names()) - unknown = [k for k in kwargs if k not in traits] - if unknown: - key = sorted(unknown)[0] - raise TypeError(f"{cls.__name__}() got unexpected keyword argument {key!r}.") - - -def _frame_to_rgb( - frame: np.ndarray, - *, - cmap: str, - vmin: float | None, - vmax: float | None, - log_scale: bool, -) -> np.ndarray: - """Colormap a 2D float frame into a uint8 (H, W, 3) RGB array.""" - values = frame.astype(np.float64, copy=False) - if log_scale: - values = np.log1p(np.clip(values - np.nanmin(values), 0, None)) - lo = float(np.nanpercentile(values, 1)) if vmin is None else float(vmin) - hi = float(np.nanpercentile(values, 99)) if vmax is None else float(vmax) - if hi <= lo: - hi = lo + 1.0 - normalized = np.clip((values - lo) / (hi - lo), 0.0, 1.0) - colormap = matplotlib.colormaps[cmap] - rgba = colormap(normalized) - return (rgba[..., :3] * 255).astype(np.uint8) - - -def _rgb_to_png_bytes(rgb: np.ndarray) -> bytes: - import io - from PIL import Image - buf = io.BytesIO() - Image.fromarray(rgb, mode="RGB").save(buf, format="PNG") - return buf.getvalue() - - class ChooseLattice(StaticFallbackMixin, anywidget.AnyWidget): """Interactive picker for an ordered origin + two lattice-vector points. @@ -155,7 +117,7 @@ def __init__( notebook_preview_max_px: int = 512, **kwargs, ) -> None: - _reject_unknown_kwargs(type(self), kwargs) + reject_unknown_kwargs(type(self), kwargs) super().__init__(**kwargs) core_image_dataset_types = _core_image_dataset_types() @@ -176,8 +138,8 @@ def __init__( ) self._data = frame - rgb = _frame_to_rgb(frame, cmap=cmap, vmin=vmin, vmax=vmax, log_scale=log_scale) - png_bytes = _rgb_to_png_bytes(rgb) + rgb = frame_to_rgb(frame, cmap=cmap, vmin=vmin, vmax=vmax, log_scale=log_scale) + png_bytes = rgb_to_png_bytes(rgb) self._configure_static_fallback( notebook_preview_format=notebook_preview_format, diff --git a/src/quantem/widget/render/__init__.py b/src/quantem/widget/render/__init__.py index 6b78a1be..31c6d6fe 100644 --- a/src/quantem/widget/render/__init__.py +++ b/src/quantem/widget/render/__init__.py @@ -1,5 +1,6 @@ """Rendering/export helpers shared by widget classes.""" +from quantem.widget.render.frame import frame_to_rgb, rgb_to_png_bytes from quantem.widget.render.thumbnail import ( save_thumbnail, thumbnail_bytes, @@ -8,6 +9,8 @@ ) __all__ = [ + "frame_to_rgb", + "rgb_to_png_bytes", "save_thumbnail", "thumbnail_bytes", "thumbnail_image", diff --git a/src/quantem/widget/render/frame.py b/src/quantem/widget/render/frame.py new file mode 100644 index 00000000..5db33728 --- /dev/null +++ b/src/quantem/widget/render/frame.py @@ -0,0 +1,59 @@ +"""Colormapping and PNG encoding for widget image panels.""" + +from __future__ import annotations + +import io + +import matplotlib +import numpy as np + +NO_DATA_COLOR = "#404040" + + +def frame_to_rgb( + frame: np.ndarray, + *, + cmap: str, + vmin: float | None = None, + vmax: float | None = None, + log_scale: bool = False, +) -> np.ndarray: + """Colormap a 2D float frame into a uint8 ``(H, W, 3)`` RGB array. + + Parameters + ---------- + frame : np.ndarray + ``(H, W)`` scalar image. + cmap : str + Matplotlib colormap name. + vmin, vmax : float, optional + Explicit display range. Defaults to a robust 1st/99th percentile + auto-contrast when not given. + log_scale : bool, default=False + Apply a ``log1p`` stretch before contrast scaling. + + Returns + ------- + np.ndarray + ``(H, W, 3)`` uint8 RGB image. + """ + values = frame.astype(np.float64, copy=False) + if log_scale: + values = np.log1p(np.clip(values - np.nanmin(values), 0, None)) + lo = float(np.nanpercentile(values, 1)) if vmin is None else float(vmin) + hi = float(np.nanpercentile(values, 99)) if vmax is None else float(vmax) + if hi <= lo: + hi = lo + 1.0 + normalized = np.clip((values - lo) / (hi - lo), 0.0, 1.0) + colormap = matplotlib.colormaps[cmap].with_extremes(bad=NO_DATA_COLOR) + rgba = colormap(normalized) + return (rgba[..., :3] * 255).astype(np.uint8) + + +def rgb_to_png_bytes(rgb: np.ndarray) -> bytes: + """Encode a uint8 ``(H, W, 3)`` RGB array as PNG bytes.""" + from PIL import Image + + buf = io.BytesIO() + Image.fromarray(rgb, mode="RGB").save(buf, format="PNG") + return buf.getvalue() diff --git a/src/quantem/widget/showbragg.py b/src/quantem/widget/showbragg.py new file mode 100644 index 00000000..9470872a --- /dev/null +++ b/src/quantem/widget/showbragg.py @@ -0,0 +1,639 @@ +"""showbragg: front end for the BraggVectors disk-detection workflow""" + +import base64 +import json +import pathlib +from typing import Any + +import anywidget +import numpy as np +import traitlets +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.datastructures.vector import Vector +from quantem.diffraction.bragg_vectors import BraggVectors +from quantem.diffraction.disk_detection import SUBPIXEL_MODES +from quantem.diffraction.strain import StrainMap + +from quantem.widget.render import frame_to_rgb, rgb_to_png_bytes +from quantem.widget.utils.array import to_numpy +from quantem.widget.utils.static_fallback import StaticFallbackMixin +from quantem.widget.utils.traits import reject_unknown_kwargs +from quantem.widget.utils.ui import UiMode, resolve_ui_mode + +TEMPLATE_SOURCES = ("synthetic", "data", "probe") +STAGE_STATES = ("idle", "preview", "running", "done", "error") + + +class ShowBragg(StaticFallbackMixin, anywidget.AnyWidget): + """Interactive driver for the BraggVectors disk-detection and lattice-fitting workflow. + + Parameters + ---------- + data : Dataset4dstem or BraggVectors + A 4D-STEM dataset, or an existing BraggVectors to reopen with its + detection kept. + probe : array_like, optional + Measured vacuum probe, required for template_source="probe". + device : str, default="cpu" + Torch device passed through to BraggVectors. + cmap : str, default="inferno" + Matplotlib colormap for every image panel. + title : str, default="" + Title above the panels. Falls back to the dataset name. + log_scale : bool, default=False + Log-stretch every image panel. + ui_mode : {"interactive", "presentation", "report", "minimal"}, default="interactive" + Display-chrome preset; explicit keywords below override it. + show_title : bool, default=True + Show the title row. + show_controls : bool, default=True + Expose the parameter controls at all. + controls_collapsed : bool, default=False + Start with the parameter controls hidden. + save_state : bool, default=False + Embed the full interactive state in the notebook. + notebook_preview_format : {"jpeg", "webp", "png"} or None, default=None + Static preview format used when save_state=False. None matches + ChooseLattice, whose live widget does not hide the fallback sibling. + notebook_preview_quality : int, default=88 + Lossy preview quality for JPEG/WebP, from 1 to 100. + notebook_preview_max_px : int, default=512 + Longest image side for the saved-notebook preview. + + Notes + ----- + template_radius, max_peak_shift and num_candidates use 0 as the sentinel for + "let the data choose"; none of the three is meaningful at zero. num_candidates + resolves to however many peaks the busiest scan position detected. + + Examples + -------- + >>> widget = ShowBragg(dataset) + >>> widget + >>> widget.detect() + >>> widget.fit() + >>> strain = widget.strain_map() + """ + + _esm = pathlib.Path(__file__).parent / "static" / "showbragg.js" + + widget_version = traitlets.Unicode("unknown").tag(sync=True) + scan_shape = traitlets.List(traitlets.Int(), default_value=[0, 0]).tag(sync=True) + q_shape = traitlets.List(traitlets.Int(), default_value=[0, 0]).tag(sync=True) + title = traitlets.Unicode("").tag(sync=True) + cmap = traitlets.Unicode("inferno").tag(sync=True) + log_scale = traitlets.Bool(False).tag(sync=True) + status = traitlets.Unicode("").tag(sync=True) + show_title = traitlets.Bool(True).tag(sync=True) + show_controls = traitlets.Bool(True).tag(sync=True) + controls_collapsed = traitlets.Bool(False).tag(sync=True) + + # Stage 1, template + template_source = traitlets.Enum(TEMPLATE_SOURCES, default_value="synthetic").tag(sync=True) + template_radius = traitlets.Float(0.0).tag(sync=True) + template_edge = traitlets.Float(1.0).tag(sync=True) + template_center = traitlets.List(traitlets.Float(), default_value=[]).tag(sync=True) + template_subtract_mean = traitlets.Bool(True).tag(sync=True) + template_roi = traitlets.List(traitlets.Float(), default_value=[]).tag(sync=True) + template_png = traitlets.Bytes(b"").tag(sync=True) + template_shape = traitlets.List(traitlets.Int(), default_value=[0, 0]).tag(sync=True) + has_probe = traitlets.Bool(False).tag(sync=True) + + # Stage 2, correlation probe + probe_position = traitlets.List(traitlets.Int(), default_value=[0, 0]).tag(sync=True) + probe_diffraction_png = traitlets.Bytes(b"").tag(sync=True) + probe_correlation_png = traitlets.Bytes(b"").tag(sync=True) + + # Stages 3 and 4, detection + min_abs_intensity = traitlets.Float(0.0).tag(sync=True) + min_spacing = traitlets.Float(0.0).tag(sync=True) + edge_boundary = traitlets.Int(1).tag(sync=True) + subpixel = traitlets.Enum(SUBPIXEL_MODES, default_value="upsample").tag(sync=True) + upsample_factor = traitlets.Int(16).tag(sync=True) + max_num_peaks = traitlets.Int(1000).tag(sync=True) + preview_grid = traitlets.Int(8).tag(sync=True) + preview_peaks = traitlets.Unicode("").tag(sync=True) + detection_state = traitlets.Enum(STAGE_STATES, default_value="idle").tag(sync=True) + + # Stages 5 and 6, Bragg vector map and basis + bvm_sampling = traitlets.Float(1.0).tag(sync=True) + bvm_png = traitlets.Bytes(b"").tag(sync=True) + num_candidates = traitlets.Int(0).tag(sync=True) + candidate_min_spacing = traitlets.Float(2.0).tag(sync=True) + candidate_min_abs_intensity = traitlets.Float(0.0).tag(sync=True) + candidates = traitlets.List(default_value=[]).tag(sync=True) + origin_index = traitlets.Int(-1).tag(sync=True) + g1_index = traitlets.Int(-1).tag(sync=True) + g2_index = traitlets.Int(-1).tag(sync=True) + origin_rc = traitlets.List(traitlets.Float(), default_value=[]).tag(sync=True) + g1_rc = traitlets.List(traitlets.Float(), default_value=[]).tag(sync=True) + g2_rc = traitlets.List(traitlets.Float(), default_value=[]).tag(sync=True) + + # Stage 7, lattice fit + min_num_peaks = traitlets.Int(5).tag(sync=True) + max_peak_shift = traitlets.Float(0.0).tag(sync=True) + fit_state = traitlets.Enum(STAGE_STATES, default_value="idle").tag(sync=True) + mask_weight_png = traitlets.Bytes(b"").tag(sync=True) + fit_error_png = traitlets.Bytes(b"").tag(sync=True) + + _TEMPLATE_TRAITS = ( + "template_source", + "template_radius", + "template_edge", + "template_center", + "template_subtract_mean", + "template_roi", + ) + _CANDIDATE_TRAITS = ( + "num_candidates", + "candidate_min_spacing", + "candidate_min_abs_intensity", + "origin_index", + "g1_index", + "g2_index", + "origin_rc", + "g1_rc", + "g2_rc", + ) + + def __init__( + self, + data, + *, + probe=None, + device: str = "cpu", + cmap: str = "inferno", + log_scale: bool = False, + title: str = "", + ui_mode: UiMode = "interactive", + show_title: bool | None = None, + show_controls: bool | None = None, + controls_collapsed: bool | None = None, + save_state: bool = False, + notebook_preview_format: str | None = None, + notebook_preview_quality: int = 88, + notebook_preview_max_px: int = 512, + **kwargs, + ) -> None: + reject_unknown_kwargs(type(self), kwargs) + super().__init__(**kwargs) + + if isinstance(data, BraggVectors): + self._bragg = data + elif isinstance(data, Dataset4dstem): + self._bragg = BraggVectors.from_dataset(data, device=device) + else: + raise TypeError( + "ShowBragg expects a Dataset4dstem or a BraggVectors instance, got " + f"{type(data).__name__}." + ) + + dataset = self._bragg.dataset + self._probe = None if probe is None else to_numpy(probe, dtype=np.float32) + + if self._bragg.template is None: + self._build_template() + + self._configure_static_fallback( + notebook_preview_format=notebook_preview_format, + notebook_preview_quality=notebook_preview_quality, + notebook_preview_max_px=notebook_preview_max_px, + ) + self._save_state = bool(save_state) + + ui = resolve_ui_mode( + ui_mode, + defaults={ + "show_title": True, + "show_controls": True, + "controls_collapsed": False, + }, + overrides={ + "show_title": show_title, + "show_controls": show_controls, + "controls_collapsed": controls_collapsed, + }, + ) + + with self.hold_sync(): + self.show_title = bool(ui["show_title"]) + self.show_controls = bool(ui["show_controls"]) + self.controls_collapsed = bool(ui["controls_collapsed"]) + self.scan_shape = [int(dataset.shape[0]), int(dataset.shape[1])] + self.q_shape = [int(dataset.shape[-2]), int(dataset.shape[-1])] + self.title = str(title or getattr(dataset, "name", "") or "") + self.cmap = str(cmap) + self.log_scale = bool(log_scale) + self.has_probe = self._probe is not None + self._render_template() + self._render_probe() + if self._bragg.peaks is not None: + self.detection_state = "done" + self._update_bvm() + self._choose_basis() + if self._bragg.u_array is not None: + self.fit_state = "done" + self._render_fit() + + self.observe(self._on_template_change, names=self._TEMPLATE_TRAITS) + self.observe(self._on_probe_change, names="probe_position") + self.observe(self._on_basis_change, names=self._CANDIDATE_TRAITS) + self.observe(self._on_display_change, names=["cmap", "log_scale"]) + self.on_msg(self._handle_msg) + + try: + from importlib.metadata import version + + self.widget_version = version("quantem-widget") + except Exception: + pass + + # Validators + + @traitlets.validate("probe_position") + def _validate_probe_position(self, proposal): + row, col = proposal["value"] + rows, cols = self.scan_shape if self.scan_shape else (1, 1) + return [ + int(np.clip(row, 0, max(0, rows - 1))), + int(np.clip(col, 0, max(0, cols - 1))), + ] + + @traitlets.validate("template_center", "origin_rc", "g1_rc", "g2_rc") + def _validate_pair(self, proposal): + value = [float(v) for v in proposal["value"]] + if value and len(value) != 2: + raise traitlets.TraitError( + f"{proposal['trait'].name} must be empty or a (row, col) pair, " + f"got {len(value)} values." + ) + return value + + @traitlets.validate("template_roi") + def _validate_roi(self, proposal): + value = [float(v) for v in proposal["value"]] + if value and len(value) != 4: + raise traitlets.TraitError( + f"template_roi must be empty or [r0, c0, r1, c1], got {len(value)} values." + ) + return value + + # Public API + + @property + def bragg(self) -> BraggVectors: + """The wrapped workflow object, for anything the panels do not expose.""" + return self._bragg + + @property + def peaks(self) -> Vector | None: + """Detected peaks, or None before a full detection run.""" + return self._bragg.peaks + + @property + def bvm(self) -> Dataset2d | None: + """Bragg vector map, or None before a full detection run.""" + return self._bragg.bvm + + @property + def basis(self) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None: + """(origin, g1, g2), with origin absolute and g1/g2 offsets from it.""" + if self._bragg.origin is None: + return None + return (self._bragg.origin, self._bragg.g1, self._bragg.g2) + + def detect(self, positions: list[tuple[int, int]] | None = None, **kwargs) -> Vector: + """Run disk detection with the current detection traits. + + Parameters + ---------- + positions : list of tuple of int, optional + (row, col) positions to test on. None runs the full scan and + populates peaks; a subset leaves the workflow state untouched. + **kwargs + Overrides for the detection traits, forwarded to detect_disks. + + Returns + ------- + Vector + Detected peaks. + """ + detect_kwargs = { + "min_abs_intensity": float(self.min_abs_intensity), + "min_spacing": float(self.min_spacing), + "edge_boundary": int(self.edge_boundary), + "subpixel": str(self.subpixel), + "upsample_factor": int(self.upsample_factor), + "max_num_peaks": int(self.max_num_peaks), + **kwargs, + } + preview = positions is not None + + self.detection_state = "preview" if preview else "running" + self.status = "" + try: + found = self._bragg.detect_disks( + positions=positions, progressbar=False, **detect_kwargs + ) + except Exception as exc: + self.detection_state = "error" + self.status = f"Detection failed: {exc}" + raise + + if preview: + self.preview_peaks = json.dumps( + { + "positions": [[int(r), int(c)] for r, c in positions], + "counts": [int(found[i].array.shape[0]) for i in range(len(positions))], + "peaks": [[float(v) for v in row] for row in found[0].array], + } + ) + self.detection_state = "done" if self._bragg.peaks is not None else "idle" + return found + + with self.hold_sync(): + self.detection_state = "done" + self._update_bvm() + self._choose_basis() + return found + + def fit(self, **kwargs) -> "ShowBragg": + """Fit the lattice at every scan position with the current fit traits. + + Parameters + ---------- + **kwargs + Overrides forwarded to fit_lattice. + + Returns + ------- + ShowBragg + self, for method chaining. + """ + fit_kwargs = { + "min_num_peaks": int(self.min_num_peaks), + "max_peak_shift": float(self.max_peak_shift) or None, + **kwargs, + } + + self.fit_state = "running" + self.status = "" + try: + self._bragg.fit_lattice(progressbar=False, plot=False, **fit_kwargs) + except Exception as exc: + self.fit_state = "error" + self.status = f"Lattice fit failed: {exc}" + raise + + with self.hold_sync(): + self.fit_state = "done" + self._render_fit() + return self + + def collapse_controls(self) -> "ShowBragg": + """Collapse the parameter controls.""" + self.controls_collapsed = True + return self + + def expand_controls(self) -> "ShowBragg": + """Expand the parameter controls.""" + self.controls_collapsed = False + return self + + def toggle_controls(self) -> "ShowBragg": + """Toggle whether the parameter controls are collapsed.""" + self.controls_collapsed = not bool(self.controls_collapsed) + return self + + def strain_map(self, **kwargs) -> StrainMap: + """Hand the fitted lattice vectors to a StrainMap. + + Parameters + ---------- + **kwargs + u_ref, v_ref and mask, forwarded to calculate_strain_map. + + Returns + ------- + StrainMap + Strain map built from the per-position lattice vectors. + """ + return self._bragg.calculate_strain_map(**kwargs) + + # Stage handlers + + def _build_template(self) -> None: + """Rebuild the correlation template from the current template traits.""" + center = tuple(self.template_center) if self.template_center else None + subtract_mean = bool(self.template_subtract_mean) + + if self.template_source == "synthetic": + self._bragg.make_template_synthetic( + radius=float(self.template_radius) or None, + edge=float(self.template_edge), + center=center, + subtract_mean=subtract_mean, + ) + elif self.template_source == "data": + self._bragg.make_template_from_data( + roi=self._roi_mask(), + subtract_mean=subtract_mean, + center=center, + ) + else: + if self._probe is None: + raise ValueError( + "template_source='probe' needs a probe image; pass " + "ShowBragg(..., probe=vacuum_probe)." + ) + self._bragg.make_template_from_probe( + self._probe, center=center, subtract_mean=subtract_mean + ) + + def _roi_mask(self) -> np.ndarray | None: + """Boolean scan mask from template_roi, or None for the whole scan.""" + if not self.template_roi: + return None + rows, cols = self.scan_shape + r0, c0, r1, c1 = self.template_roi + mask = np.zeros((rows, cols), dtype=bool) + mask[ + int(np.clip(min(r0, r1), 0, rows)) : int(np.clip(max(r0, r1), 0, rows)) + 1, + int(np.clip(min(c0, c1), 0, cols)) : int(np.clip(max(c0, c1), 0, cols)) + 1, + ] = True + return mask + + def _preview_positions(self) -> list[tuple[int, int]]: + """Probe position first, then an evenly spaced grid across the scan.""" + rows, cols = self.scan_shape + n = max(1, int(self.preview_grid)) + grid_rows = np.unique(np.linspace(0, rows - 1, min(n, rows)).astype(int)) + grid_cols = np.unique(np.linspace(0, cols - 1, min(n, cols)).astype(int)) + + probe = (int(self.probe_position[0]), int(self.probe_position[1])) + positions = [probe] + positions += [ + (int(r), int(c)) for r in grid_rows for c in grid_cols if (int(r), int(c)) != probe + ] + return positions + + def _candidate_count(self) -> int: + """How many candidate peaks to number, from the data unless set explicitly.""" + if self.num_candidates > 0: + return int(self.num_candidates) + + peaks = self._bragg.peaks + counts = peaks.row_counts() if peaks is not None else [] + return max(4, int(max(counts))) if len(counts) else 100 + + def _choose_basis(self) -> None: + """Re-derive the candidate peaks and resolve the basis from the picks.""" + if self._bragg.bvm is None: + return + + try: + self._bragg.choose_basis_vectors( + origin=self._basis_pick("origin"), + g1=self._basis_pick("g1"), + g2=self._basis_pick("g2"), + num_candidates=self._candidate_count(), + min_spacing=float(self.candidate_min_spacing), + min_abs_intensity=float(self.candidate_min_abs_intensity), + plot=False, + ) + self._bragg.index_peaks(plot=False) + except Exception as exc: + self.status = f"Basis selection failed: {exc}" + return + + cand_rc = self._bragg.candidates_rc + cand_int = self._bragg.candidates_intensity + + self._resolving_basis = True + try: + with self.hold_sync(): + self.status = "" + self.candidates = [ + [float(rc[0]), float(rc[1]), float(i)] + for rc, i in zip(cand_rc, cand_int) + ] + self.origin_rc = [float(v) for v in self._bragg.origin] + self.g1_rc = [float(v) for v in self._bragg.g1] + self.g2_rc = [float(v) for v in self._bragg.g2] + finally: + self._resolving_basis = False + + def _basis_pick(self, name: str) -> int | list[float] | None: + """One basis argument: candidate index, explicit vector, or None for auto.""" + index = int(getattr(self, f"{name}_index")) + if index >= 0: + return index + vector = getattr(self, f"{name}_rc") + return list(vector) if vector else None + + # Rendering + + def _render(self, frame) -> bytes: + """Colormap a 2D array into PNG bytes, or empty bytes when absent.""" + if frame is None: + return b"" + return rgb_to_png_bytes( + frame_to_rgb( + np.asarray(frame, dtype=float), + cmap=self.cmap, + log_scale=bool(self.log_scale), + ) + ) + + def _render_template(self) -> None: + """Render the template cropped to the disk, which is tiny on a large detector.""" + template = self._bragg.template + if template is None: + self.template_png = b"" + return + + frame = np.asarray(template, dtype=float) + rows, cols = frame.shape + centre = (rows // 2, cols // 2) + + strong = np.argwhere(frame > 0.25 * frame.max()) + reach = 8.0 if strong.size == 0 else np.abs(strong - np.array(centre)).max() + half = int(max(12, min(min(rows, cols) // 2, round(3 * reach)))) + + crop = frame[ + centre[0] - half : centre[0] + half + 1, + centre[1] - half : centre[1] + half + 1, + ] + with self.hold_sync(): + self.template_shape = [int(crop.shape[0]), int(crop.shape[1])] + self.template_png = self._render(crop) + + def _render_probe(self) -> None: + row, col = self.probe_position + pattern = np.asarray(self._bragg.dataset.array[row, col]) + with self.hold_sync(): + self.probe_diffraction_png = self._render(pattern) + self.probe_correlation_png = self._render(self._bragg.correlation_map(row, col)) + + def _update_bvm(self) -> None: + bvm = self._bragg.compute_bvm(float(self.bvm_sampling)) + self.bvm_png = self._render(bvm.array) + + def _render_fit(self) -> None: + with self.hold_sync(): + self.mask_weight_png = self._render(self._bragg.mask_weight) + self.fit_error_png = self._render(self._bragg.fit_error) + + # Observers and messages + + def _on_template_change(self, _change) -> None: + try: + self._build_template() + except Exception as exc: + self.status = str(exc) + return + with self.hold_sync(): + self.status = "" + self._render_template() + self._render_probe() + + def _on_probe_change(self, _change) -> None: + self._render_probe() + + def _on_basis_change(self, _change) -> None: + if getattr(self, "_resolving_basis", False): + return + self._choose_basis() + + def _on_display_change(self, _change) -> None: + with self.hold_sync(): + self._render_template() + self._render_probe() + if self._bragg.bvm is not None: + self.bvm_png = self._render(self._bragg.bvm.array) + if self._bragg.u_array is not None: + self._render_fit() + + def _handle_msg(self, _widget: Any, content: dict[str, Any], _buffers: list[Any]) -> None: + """Panel actions. Errors land in status instead of the kernel.""" + action = content.get("type") + try: + if action == "preview_detect": + self.detect(positions=self._preview_positions()) + elif action == "run_detect": + self.detect() + elif action == "compute_bvm": + self._update_bvm() + self._choose_basis() + elif action == "run_fit": + self.fit() + except Exception as exc: + if not self.status: + self.status = str(exc) + + def _static_png_b64(self, max_px: int = 512) -> str | None: + preview = self.bvm_png or self.probe_correlation_png or self.template_png + if not preview: + return None + return base64.b64encode(bytes(preview)).decode("ascii") diff --git a/src/quantem/widget/utils/traits.py b/src/quantem/widget/utils/traits.py new file mode 100644 index 00000000..841e29d3 --- /dev/null +++ b/src/quantem/widget/utils/traits.py @@ -0,0 +1,12 @@ +"""Trait-related helpers shared by widget constructors.""" + +from __future__ import annotations + + +def reject_unknown_kwargs(cls, kwargs: dict) -> None: + """Raise TypeError for any kwarg that isn't a declared trait (catches typos).""" + traits = set(cls.class_trait_names()) + unknown = [k for k in kwargs if k not in traits] + if unknown: + key = sorted(unknown)[0] + raise TypeError(f"{cls.__name__}() got unexpected keyword argument {key!r}.") diff --git a/tests/showbragg/test_showbragg.py b/tests/showbragg/test_showbragg.py new file mode 100644 index 00000000..69c4d2da --- /dev/null +++ b/tests/showbragg/test_showbragg.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import io + +import numpy as np +import pytest +from PIL import Image + +from quantem.widget.showbragg import ShowBragg + +pytest.importorskip("quantem.diffraction.bragg_vectors") + +from quantem.core.datastructures.dataset4dstem import Dataset4dstem # noqa: E402 +from quantem.diffraction.bragg_vectors import BraggVectors # noqa: E402 +from quantem.diffraction.strain import StrainMap # noqa: E402 + +G1 = np.array([0.0, 7.0]) +G2 = np.array([7.0, 0.0]) + + +def synthetic_dataset(scan=(6, 5), detector=(32, 32)) -> Dataset4dstem: + """Square lattice of soft disks, identical at every scan position.""" + rows, cols = scan + height, width = detector + rr, cc = np.mgrid[0:height, 0:width] + origin = np.array([height / 2, width / 2]) + + pattern = np.zeros((height, width), dtype=np.float32) + for a in (-1, 0, 1): + for b in (-1, 0, 1): + spot = origin + a * G1 + b * G2 + pattern += np.exp(-((rr - spot[0]) ** 2 + (cc - spot[1]) ** 2) / 4.0) + + data = np.broadcast_to(pattern, (rows, cols, height, width)).copy() + return Dataset4dstem.from_array(data, name="synthetic lattice") + + +@pytest.fixture +def widget() -> ShowBragg: + return ShowBragg(synthetic_dataset()) + + +@pytest.fixture +def fitted(widget: ShowBragg) -> ShowBragg: + widget.detect() + widget.fit() + return widget + + +def test_construction_reports_shapes_and_renders_template(widget: ShowBragg) -> None: + assert widget.scan_shape == [6, 5] + assert widget.q_shape == [32, 32] + assert widget.title == "synthetic lattice" + assert widget.peaks is None + + decoded = Image.open(io.BytesIO(bytes(widget.template_png))).convert("RGB") + assert decoded.size == (widget.template_shape[1], widget.template_shape[0]) + assert 0 < widget.template_shape[0] <= widget.q_shape[0] + + +def test_construction_from_bragg_vectors_keeps_detected_state() -> None: + bragg = BraggVectors.from_dataset(synthetic_dataset()) + bragg.make_template_synthetic(radius=2.0) + bragg.detect_disks(progressbar=False) + + widget = ShowBragg(bragg) + + assert widget.bragg is bragg + assert widget.peaks is not None + assert widget.bvm is not None + assert len(widget.bvm_png) > 0 + + +def test_rejects_bad_data_and_unknown_kwarg() -> None: + with pytest.raises(TypeError): + ShowBragg(np.random.rand(16, 16).astype(np.float32)) + with pytest.raises(TypeError): + ShowBragg(synthetic_dataset(), not_a_real_kwarg=True) + + +def test_template_trait_change_rerenders_preview(widget: ShowBragg) -> None: + before = bytes(widget.template_png) + widget.template_radius = 5.0 + assert bytes(widget.template_png) != before + + +def test_probe_position_clamped_to_scan_shape(widget: ShowBragg) -> None: + widget.probe_position = [99, -4] + assert widget.probe_position == [5, 0] + + +def test_preview_detection_leaves_full_dataset_peaks_empty(widget: ShowBragg) -> None: + preview = widget.detect(positions=[(0, 0), (2, 3)]) + + assert preview.shape == (2,) + assert widget.peaks is None + assert widget.bvm is None + assert widget.preview_peaks != "" + + +def test_basis_accepts_candidate_index_or_vector(widget: ShowBragg) -> None: + widget.detect() + + widget.origin_index = 0 + widget.g1_index = 1 + widget.g2_index = 2 + by_index = tuple(np.array(v) for v in widget.basis) + + widget.origin_index = -1 + widget.g1_index = -1 + widget.g2_index = -1 + widget.origin_rc = list(by_index[0]) + widget.g1_rc = list(by_index[1]) + widget.g2_rc = list(by_index[2]) + by_vector = tuple(np.array(v) for v in widget.basis) + + for expected, actual in zip(by_index, by_vector): + np.testing.assert_allclose(actual, expected) + + origin = np.asarray(widget.candidates[0][:2], dtype=float) + np.testing.assert_allclose(by_index[0], origin) + + +def test_candidate_count_defaults_to_the_busiest_scan_position(widget: ShowBragg) -> None: + assert widget.num_candidates == 0 + + widget.detect() + busiest = max(widget.peaks.row_counts()) + assert len(widget.candidates) <= busiest + + widget.num_candidates = 5 + assert len(widget.candidates) <= 5 + + +def test_detect_and_fit_populate_results(fitted: ShowBragg) -> None: + assert fitted.detection_state == "done" + assert fitted.fit_state == "done" + assert fitted.peaks.shape == (6, 5) + assert fitted.bragg.u_array.shape == (6, 5, 2) + assert len(fitted.mask_weight_png) > 0 + assert len(fitted.fit_error_png) > 0 + + +def test_ui_mode_preset_and_control_toggles() -> None: + widget = ShowBragg(synthetic_dataset(), ui_mode="report") + assert widget.show_title is True + assert widget.show_controls is False + + widget.expand_controls() + assert widget.controls_collapsed is False + widget.toggle_controls() + assert widget.controls_collapsed is True + + explicit = ShowBragg(synthetic_dataset(), ui_mode="minimal", show_title=True) + assert explicit.show_title is True + assert explicit.show_controls is False + + +def test_strain_map_requires_a_fit(widget: ShowBragg) -> None: + with pytest.raises(ValueError, match="fit_lattice"): + widget.strain_map() + + widget.detect() + widget.fit() + assert isinstance(widget.strain_map(), StrainMap)