diff --git a/.github/workflows/widget-release.yml b/.github/workflows/widget-release.yml index 9b60f61f..660b6185 100644 --- a/.github/workflows/widget-release.yml +++ b/.github/workflows/widget-release.yml @@ -130,6 +130,7 @@ jobs: required = { "quantem/widget/static/chooselattice.js", "quantem/widget/static/show1d.js", + "quantem/widget/static/plot2d.js", "quantem/widget/static/show2d.js", "quantem/widget/static/show3d.js", "quantem/widget/static/show3dslices.js", diff --git a/CHANGELOG.md b/CHANGELOG.md index af09ba71..b0053fed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,9 @@ new `rcN` heading when that rc is published to TestPyPI. ## Unreleased +- Add `Plot2D` for scalar maps with calibrated Cartesian axes, colormap + selection, viewport controls, and editable Matplotlib figure export. + - Maintainer docs split pull requests into discuss-first (new widgets, cross-widget refactors) and incremental in-widget fixes, and add a `widget-tutorials/` upload page for the public diff --git a/README.md b/README.md index fb729804..ce95f115 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,7 @@ for backend setup, Colab instructions, and verification. | Widget | Use it for | Learn more | |---|---|---| +| `Plot2D` | Scalar maps with physical axes, color scales, and calibrated hover | [tutorial](docs/tutorials/plot2d.ipynb) · [API](docs/api/plot2d.md) | | `Show1D` | Scientific traces, reconstruction metrics, and live monitors | [API](https://electronmicroscopy.github.io/quantem.widget/api/show1d.html) | | `Show2D` | Images, contrast, FFTs, ROIs, profiles, and scale bars | [tutorial](https://electronmicroscopy.github.io/quantem.widget/tutorials/show2d.html) · [API](https://electronmicroscopy.github.io/quantem.widget/api/show2d.html) | | `Mask2D` | Draw one rectangle, square, or circle and use its Boolean mask directly in Python | [guide and API](https://electronmicroscopy.github.io/quantem.widget/api/mask2d.html) | @@ -53,6 +54,16 @@ for backend setup, Colab instructions, and verification. ## Documentation +For a scalar map, import `Plot2D` from the same package: + +```python +from quantem.widget import Plot2D + +# values.shape == (len(angle), len(radius)); coordinates are bin centers. +plot = Plot2D(values, x=radius, y=angle, + x_label="Distance (Å)", y_label="Angle (°)") +``` + Visit the **[quantem.widget documentation](https://electronmicroscopy.github.io/quantem.widget/)** for installation, tutorials, API references, command-line workflows, data I/O, HTML sharing, and WebGPU export guidance. diff --git a/docs/_toc.yml b/docs/_toc.yml index 2c051c44..bbc81c7d 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -24,6 +24,7 @@ parts: title: Compare datasets or tilts - file: tutorials/showptycho - file: tutorials/show1d + - file: tutorials/plot2d - file: tutorials/show2d - file: tutorials/show3d - file: tutorials/show3dslices @@ -67,6 +68,7 @@ parts: - file: api/datasets - file: api/viewer-ui - file: api/show1d + - file: api/plot2d - file: api/show2d - file: api/mask2d - file: api/show3d @@ -117,6 +119,8 @@ parts: sections: - file: maintainer/storyboard-show2d title: Show2D + - file: maintainer/storyboard-plot2d + title: Plot2D - file: maintainer/storyboard-show3d title: Show3D - file: maintainer/storyboard-show3dslices diff --git a/docs/api/index.md b/docs/api/index.md index 3c030a2b..ce730e0e 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -42,6 +42,7 @@ readers and test agents. | Widget | Class | Offline export | |---|---|---| +| [Plot2D](plot2d) | `quantem.widget.Plot2D` | PNG from browser; Matplotlib figures via Python; no standalone HTML API | | [Show1D](show1d) | `quantem.widget.show1d.Show1D` | state JSON, CSV, PNG/PDF via Python, interactive HTML | | [Show2D](show2d) | `quantem.widget.show2d.Show2D` | state JSON, PNG, interactive HTML (`encoding="full"` / `encoding="uint8"`) | | [Mask2D](mask2d) | `quantem.widget.mask2d.Mask2D` | Boolean mask and optional selected geometry in Python | diff --git a/docs/api/plot2d.md b/docs/api/plot2d.md new file mode 100644 index 00000000..d655d107 --- /dev/null +++ b/docs/api/plot2d.md @@ -0,0 +1,77 @@ +# Plot2D + +See the [interactive tutorial](../tutorials/plot2d.ipynb). + +Use `quantem.widget.Plot2D` for a calibrated scalar map, such as G3 versus +distance and angle. Use Show2D for spatial images. Plot2D owns axes, colorbar, +zoom/pan and hover; it does not calculate correlations or train a model. + +```python +import numpy as np +import quantem.widget as qw + +qw.profile(check_updates=False) +radius = (np.arange(100) + 0.5) * 0.1 +angle = (np.arange(36) + 0.5) * 5 +values = np.cos(np.deg2rad(angle[:, None])) ** 2 * radius[None, :] +plot = qw.Plot2D( + values, x=radius, y=angle, + x_label="Second-neighbor distance (Å)", y_label="Shared-root angle (°)", + colorbar_label="Illustrative value", width=500, max_width=600, +) +plot +``` + +`x` contains column bin centers; `y` contains row bin centers. Both must be +increasing uniform grids. Row zero is at the bottom, matching Cartesian plots. +Source values and browser hover transport retain float64. The shared QuantEM +WebGPU colormap renderer uses float32 display buffers. Canvas fallback is +explicitly labeled when WebGPU is unavailable; it is not GPU acceleration. + +Wheel over the map to zoom, then drag to pan. Zoom buttons and Reset View are +also available. The Color menu changes map and colorbar together; each plot is +independent. A browser-local animation-frame scheduler handles gestures without +Python round trips. Stable view bounds are saved after interaction. + +Map pixels, hover values and color-scale metadata update together after +rendering completes, including during rapid replacements. + +`plot.set_data(next_values)` preserves the original grid, color limits and +viewport. `plot.horizontal_line = 92.5` adds an angle-reading line without +resending the map. `plot.figure()` returns a closed Matplotlib figure with +the current axes, values, colormap and viewport, for example +`plot.figure().savefig("g3.svg")`. By default, saved snapshots omit the map array and retain a static PNG +preview. Use `save_state=True` to embed the complete float64 map for interactive +restoration in supporting frontends. Static previews record the view when +Python renders the preview or creates a full snapshot; they do not track +browser-only gestures continuously. Notebook-manager save/restore behavior +varies by frontend. Keep scientific data files separately. + +## Current scope + +This API targets small, finite scalar maps, not large spatial images. It does +not support nonuniform coordinates, logarithmic axes, or standalone +`export_html`. Full interactive embedding is opt-in with `save_state=True`; +keep large data outside notebooks. Static previews use stride sampling above +512 bins per axis, preserve calibrated bounds, and are for viewing only. Controls, axes and the canvas follow the +notebook or documentation light/dark theme. + +## Reference + +```{eval-rst} +.. autoclass:: quantem.widget.Plot2D + :members: set_data, figure +``` + +## Interactive controls + +| Control | Behavior | +|---|---| +| Color | Recolor the map and scale without modifying values. | +| Zoom In / Zoom Out | Zoom about the viewport center. | +| Wheel / drag | Zoom about the pointer; pan within the full grid. | +| Reset View / double-click | Restore full physical bounds. | +| Save PNG | Save the current canvas, including labels and color scale. | +| Hover | Inspect original `(row, col)`, calibrated coordinates and value. | + +The [storyboard](../maintainer/storyboard-plot2d.md) defines browser signoff. diff --git a/docs/maintainer/storyboard-plot2d.md b/docs/maintainer/storyboard-plot2d.md new file mode 100644 index 00000000..23ba70ed --- /dev/null +++ b/docs/maintainer/storyboard-plot2d.md @@ -0,0 +1,20 @@ +# Plot2D storyboard + +Use the [Plot2D tutorial](../tutorials/plot2d.ipynb) for the deterministic +calibration fixture. Repeat with a representative measured or computed scalar +map for scientific signoff. These are required stories, not completed results. + +| ID | Action | Required evidence | +|---|---|---| +| P2D-01 | Hover low/high row and column bins | Physical bin centers and original float64 values match the input; row zero is at the bottom. | +| P2D-02 | Wheel zoom in/out, drag, then Reset View | Scientific pixels and axis limits change together; reset restores the full grid; source values stay unchanged. | +| P2D-03 | Change Color on one of two plots | Map and colorbar change together, numerical limits stay fixed, and the other plot stays unchanged. | +| P2D-04 | Move the tutorial angle slider; replace values with set_data | Reading line moves without resending data; replacement map updates without changing viewport or color limits. | +| P2D-05 | Use narrow/wide layouts, light/dark notebook themes | Labels, menus and axes remain readable; max_width is respected; no clipped controls. | +| P2D-06 | Save PNG, export a Matplotlib figure | Both exports show the current axes, values, viewport, color limits and reading line. | +| P2D-07 | Save notebook after interaction, close and reopen | Test default static preview and opt-in save_state=True separately. Default snapshot omits data_bytes; opt-in preserves float64 arrays. Saved view restores without a kernel where supported; payload size and fallback freshness are stated. | +| P2D-08 | Rapidly replace maps/colormaps while hovering; close during rendering | Hover corresponds to the visible map; stale asynchronous work cannot paint later; resources are released without errors. | + +Record browser/adapter, map shape and dtype, screenshots before/after, console +errors, first-paint time and gesture-to-paint latency. Python state tests and +cell execution alone are not browser signoff. diff --git a/docs/tutorials/plot2d.ipynb b/docs/tutorials/plot2d.ipynb new file mode 100644 index 00000000..914593b3 --- /dev/null +++ b/docs/tutorials/plot2d.ipynb @@ -0,0 +1,167 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Inspect a calibrated scalar map\n", + "\n", + "Use Plot2D when each array axis is a scientific quantity, such as distance and angle, rather than image position. This small synthetic map teaches display and selection; it is not a physical G3 calculation or a model prediction.\n", + "\n", + "Run using a development build containing Plot2D. No GPU training, real-data download, or learned weights are required." + ], + "id": "plot2d-00" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import ipywidgets as widgets\n", + "import quantem.widget as qw\n", + "\n", + "qw.profile(check_updates=False)" + ], + "id": "plot2d-01" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Make a known map\n", + "\n", + "Columns are radius bin centers in Å; rows are angle bin centers in degrees. Row zero is shown at the bottom. The illustrative peak is near 2.4 Å and 110°. The 36 × 60 float64 map occupies about 17 kB before widget metadata." + ], + "id": "plot2d-02" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ], + "mystnb": { + "code_prompt_show": "Show illustrative map generation" + } + }, + "outputs": [], + "source": [ + "radius = (np.arange(60) + 0.5) * 0.1\n", + "angle = (np.arange(36) + 0.5) * 5.0\n", + "values = np.exp(-((radius[None, :] - 2.4) / 0.35) ** 2\n", + " - ((angle[:, None] - 110.0) / 15.0) ** 2)" + ], + "id": "plot2d-03" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Inspect with physical axes\n", + "\n", + "Hover near the peak and read its bin, radius, angle and value. Wheel to zoom, drag to pan, and select Reset View. Changing Color changes only the display. The source and hover values remain float64; browser colormapping uses float32 display buffers." + ], + "id": "plot2d-04" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot = qw.Plot2D(\n", + " values, x=radius, y=angle,\n", + " x_label=\"Distance (Å)\", y_label=\"Angle (°)\",\n", + " colorbar_label=\"Illustrative value\", title=\"Calibrated map\",\n", + " width=500, max_width=600, vmin=0, vmax=1, save_state=True,\n", + ")\n", + "plot" + ], + "id": "plot2d-05" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Move a reading line\n", + "\n", + "The slider selects an angle row; it does not change the data. The link is browser-local and does not require Python callbacks during dragging." + ], + "id": "plot2d-06" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "angle_slider = widgets.FloatSlider(\n", + " value=107.5, min=2.5, max=177.5, step=5.0,\n", + " description=\"Angle (°)\", continuous_update=True,\n", + ")\n", + "plot.horizontal_line = angle_slider.value\n", + "angle_link = widgets.jslink((angle_slider, \"value\"), (plot, \"horizontal_line\"))\n", + "angle_slider" + ], + "id": "plot2d-07" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Replace values, keep the view\n", + "\n", + "Zoom first, then run the next cell. The same plot updates in place, retaining the grid, color limits and viewport. A second Plot2D display is not created." + ], + "id": "plot2d-08" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot.set_data(values * 0.75)" + ], + "id": "plot2d-09" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Obtain an editable figure\n", + "\n", + "The figure retains the current data, labels and viewport. This is a deliberate static export preview, separate from the interactive plot; it is displayed once. To save it, call `figure.savefig(\"scalar-map.svg\")`.\n", + "\n", + "The widget follows the notebook light/dark theme. By default, saved snapshots omit the array and keep a static PNG preview. This small tutorial explicitly uses `save_state=True` so its complete float64 map remains interactive in supporting notebook and documentation frontends. Retain scientific data separately. Standalone HTML export is not part of this API." + ], + "id": "plot2d-10" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "figure = plot.figure()\n", + "figure" + ], + "id": "plot2d-11" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/js/plot2d/index.tsx b/js/plot2d/index.tsx new file mode 100644 index 00000000..7f64c871 --- /dev/null +++ b/js/plot2d/index.tsx @@ -0,0 +1,539 @@ +import * as React from "react"; +import { createRoot } from "react-dom/client"; +import Select from "@mui/material/Select"; +import MenuItem from "@mui/material/MenuItem"; +import { + COLORMAPS, + createGPUColormapEngine, + GPUColormapEngine, + renderToOffscreen, +} from "../colormaps"; +import { extractBytes, downloadBlob } from "../format"; +import { detectTheme, getThemeColors, useTheme } from "../theme"; +import { useHideStaticFallback } from "../staticFallback"; + +type Model = { + get(key: string): any; + set(key: string, value: unknown): void; + save_changes(): void; + on(event: string, callback: () => void): void; + off(event: string, callback: () => void): void; +}; + +function render({ model, el }: { model: Model; el: HTMLElement }) { + let themeColors = getThemeColors(detectTheme().theme); + const host = document.createElement("div"); + host.dataset.quantemPlot2d = "true"; + host.style.cssText = + "width:100%;min-width:260px;font:12px system-ui;"; + const preview = document.createElement("img"); + preview.alt = "Plot2D saved preview; rerun the cell for interaction"; + preview.style.cssText = "display:none;width:100%;height:auto"; + const canvas = document.createElement("canvas"); + canvas.dataset.quantemScientificOutput = "plot2d-map"; + canvas.style.cssText = + "width:100%;display:block;touch-action:none;cursor:crosshair"; + const controls = document.createElement("div"); + controls.style.cssText = + "display:flex;gap:8px;align-items:center;flex-wrap:wrap;padding:4px 8px"; + const reset = document.createElement("button"); + reset.textContent = "Reset View"; + const zoomIn = document.createElement("button"); + zoomIn.textContent = "Zoom In"; + const zoomOut = document.createElement("button"); + zoomOut.textContent = "Zoom Out"; + const zoomLabel = document.createElement("span"); + zoomLabel.setAttribute("aria-live", "polite"); + const save = document.createElement("button"); + save.textContent = "Save PNG"; + const status = document.createElement("span"); + status.textContent = "Preparing display…"; + const readout = document.createElement("div"); + readout.style.cssText = + "height:22px;padding:3px 8px;font-variant-numeric:tabular-nums;white-space:nowrap;overflow:hidden;text-overflow:ellipsis"; + const colorControl = document.createElement("span"); + colorControl.style.cssText = "display:inline-flex;gap:6px;align-items:center;flex-shrink:0"; + const colorRoot = createRoot(colorControl); + function ColorControl() { + const { colors } = useTheme(); + useHideStaticFallback(model, { current: host }, Boolean(bitmap)); + React.useLayoutEffect(() => { + themeColors = colors; + host.style.color = colors.text; + host.style.background = colors.bg; + status.style.color = zoomLabel.style.color = colors.textMuted; + for (const button of [reset, zoomIn, zoomOut, save]) { + Object.assign(button.style, { + color: colors.text, background: colors.controlBg, + border: `1px solid ${colors.border}`, borderRadius: "3px", + font: "10px system-ui", padding: "3px 8px", cursor: "pointer", + }); + } + schedule(); + }, [colors]); + return <> + Color + + ; + } + function updateColorControl() { + colorRoot.render(); + } + controls.append(colorControl, zoomIn, zoomOut, save, reset, zoomLabel, status); + host.append(preview, canvas, controls, readout); + el.append(host); + let engine: GPUColormapEngine | null = null; + let bitmap: CanvasImageSource | null = null; + let source = new Float64Array(); + let displayCmap = model.get("cmap"); + let displayMin = model.get("vmin"); + let displayMax = model.get("vmax"); + let pointer: { clientX: number; clientY: number } | null = null; + let disposed = false, + generation = 0, + frame = 0; + let queue = Promise.resolve(); + let bounds = ( + model.get("view_bounds").length + ? model.get("view_bounds") + : model.get("grid").bounds + ).slice(); + let drag: { x: number; y: number; bounds: number[] } | null = null; + let wheelTimer = 0; + let readoutFrame = 0; + function showReadout() { + if (!readoutFrame) readoutFrame = requestAnimationFrame(() => { + readoutFrame = 0; + updateReadout(); + }); + } + const full = () => model.get("grid").bounds as number[]; + function size() { + host.style.maxWidth = `${model.get("max_width") ?? 600}px`; + schedule(); + } + const geometry = () => ({ + width: Math.max(260, host.clientWidth), + height: model.get("plot_height_px"), + left: 65, + top: 30, + right: 18, + bottom: 105, + }); + let paintedBounds = bounds.slice(); + let paintedGeometry = geometry(); + const number = (value: number) => Number(value.toPrecision(4)).toString(); + const tickNumber = (value: number, span: number) => + number(Math.abs(value) < span * 1e-12 ? 0 : value); + function commit() { + model.set("view_bounds", bounds.slice()); + model.save_changes(); + } + function clampBounds(next: number[]) { + const original = full(); + return [0, 2].flatMap((index) => { + const span = Math.min( + original[index + 1] - original[index], + Math.max( + (original[index + 1] - original[index]) / 100, + next[index + 1] - next[index], + ), + ); + const low = Math.max( + original[index], + Math.min(original[index + 1] - span, next[index]), + ); + return [low, low + span]; + }); + } + function paint() { + frame = 0; + if (disposed || canvas.style.display === "none") return; + const g = geometry(), + width = g.width - g.left - g.right, + height = g.height - g.top - g.bottom; + const ratio = window.devicePixelRatio || 1; + canvas.width = Math.round(g.width * ratio); + canvas.height = Math.round(g.height * ratio); + canvas.style.height = `${g.height}px`; + const ctx = canvas.getContext("2d")!; + ctx.scale(ratio, ratio); + ctx.fillStyle = themeColors.bg; + ctx.fillRect(0, 0, g.width, g.height); + const grid = model.get("grid"), + original = full(); + const zoom = (original[1] - original[0]) / (bounds[1] - bounds[0]); + zoomLabel.textContent = `${number(zoom)}× · ${zoom > 1.001 ? "Drag to pan" : "Zoom in to pan"}`; + canvas.style.cursor = drag + ? "grabbing" + : zoom > 1.001 + ? "grab" + : "crosshair"; + if (bitmap) { + ctx.save(); + ctx.translate(g.left, g.top + height); + ctx.scale(1, -1); + ctx.imageSmoothingEnabled = false; + ctx.drawImage( + bitmap, + ((bounds[0] - original[0]) / (original[1] - original[0])) * grid.cols, + ((bounds[2] - original[2]) / (original[3] - original[2])) * grid.rows, + ((bounds[1] - bounds[0]) / (original[1] - original[0])) * grid.cols, + ((bounds[3] - bounds[2]) / (original[3] - original[2])) * grid.rows, + 0, + 0, + width, + height, + ); + ctx.restore(); + } + ctx.strokeStyle = themeColors.border; + ctx.strokeRect(g.left, g.top, width, height); + ctx.font = "12px system-ui"; + ctx.fillStyle = themeColors.text; + ctx.textAlign = "center"; + ctx.fillText(model.get("title"), g.left + width / 2, 17); + for (let tick = 0; tick <= 4; tick++) { + const fraction = tick / 4; + ctx.textAlign = "center"; + ctx.fillText( + tickNumber( + bounds[0] + fraction * (bounds[1] - bounds[0]), + bounds[1] - bounds[0], + ), + g.left + fraction * width, + g.top + height + 17, + ); + ctx.textAlign = "right"; + ctx.fillText( + tickNumber( + bounds[2] + fraction * (bounds[3] - bounds[2]), + bounds[3] - bounds[2], + ), + g.left - 7, + g.top + height * (1 - fraction) + 4, + ); + } + ctx.textAlign = "center"; + ctx.fillText(model.get("x_label"), g.left + width / 2, g.top + height + 37); + ctx.save(); + ctx.translate(15, g.top + height / 2); + ctx.rotate(-Math.PI / 2); + ctx.fillText(model.get("y_label"), 0, 0); + ctx.restore(); + const lut = COLORMAPS[displayCmap]; + if (lut) + for (let i = 0; i < 256; i++) { + ctx.fillStyle = `rgb(${lut[i * 3]},${lut[i * 3 + 1]},${lut[i * 3 + 2]})`; + ctx.fillRect( + g.left + (i * width) / 256, + g.height - 52, + width / 256 + 0.5, + 10, + ); + } + ctx.fillStyle = themeColors.text; + ctx.textAlign = "left"; + ctx.fillText(number(displayMin), g.left, g.height - 27); + ctx.textAlign = "right"; + ctx.fillText(number(displayMax), g.left + width, g.height - 27); + ctx.textAlign = "center"; + ctx.fillText(model.get("colorbar_label"), g.left + width / 2, g.height - 8); + const line = model.get("horizontal_line"); + if (line != null && line >= bounds[2] && line <= bounds[3]) { + const pos = + g.top + height * (1 - (line - bounds[2]) / (bounds[3] - bounds[2])); + ctx.strokeStyle = "#e649a0"; + ctx.setLineDash([3, 3]); + ctx.beginPath(); + ctx.moveTo(g.left, pos); + ctx.lineTo(g.left + width, pos); + ctx.stroke(); + } + paintedBounds = bounds.slice(); + paintedGeometry = g; + updateReadout(); + canvas.setAttribute( + "aria-label", + `${model.get("title")}; ${model.get("x_label")}; ${model.get("y_label")}; ${model.get("colorbar_label")}`, + ); + canvas.dataset.paintGeneration = String( + Number(canvas.dataset.paintGeneration || 0) + 1, + ); + } + function schedule() { + if (!frame) frame = requestAnimationFrame(paint); + } + const ready = createGPUColormapEngine() + .then((value) => { + engine = value; + return value; + }) + .catch(() => null); + function prepare() { + const current = ++generation; + const bytes = extractBytes(model.get("data_bytes")); + if (!bytes.length) { + const saved = model.get("_static_fallback_jpeg"); + if (saved) preview.src = `data:${model.get("_static_fallback_mime") || "image/png"};base64,${saved}`; + else preview.removeAttribute("src"); + preview.style.display = saved ? "block" : "none"; + canvas.style.display = controls.style.display = "none"; + readout.textContent = "Saved preview · rerun the cell for interaction"; + return; + } + queue = queue + .then(async () => { + await ready; + if (disposed || current !== generation) return; + const grid = model.get("grid"), + bytes = extractBytes(model.get("data_bytes")); + const nextSource = new Float64Array( + bytes.slice(0, grid.rows * grid.cols * 8).buffer, + ); + const cmap = model.get("cmap"), + vmin = model.get("vmin"), vmax = model.get("vmax"); + const display = Float32Array.from(nextSource), lut = COLORMAPS[cmap]; + if (!lut) { + status.textContent = "Unsupported colormap"; + return; + } + let next: CanvasImageSource | null = null; + if (engine) { + engine.uploadData(0, display, grid.cols, grid.rows); + engine.uploadLUT(cmap, lut); + const rendered = await engine.renderSlotsToImageBitmapAsync( + [0], + [{ vmin, vmax }], + ); + next = rendered?.[0] ?? null; + } + const usedGPU = Boolean(next); + if (!next) { + next = renderToOffscreen( + display, + grid.cols, + grid.rows, + lut, + vmin, + vmax, + ); + } + if (disposed || current !== generation) { + if (next instanceof ImageBitmap) next.close(); + return; + } + if (bitmap instanceof ImageBitmap) bitmap.close(); + // Publish source, color metadata and pixels in the same synchronous paint. + bitmap = next; + preview.style.display = "none"; + canvas.style.display = "block"; + controls.style.display = "flex"; + source = nextSource; + displayCmap = cmap; + displayMin = vmin; + displayMax = vmax; + status.textContent = `${usedGPU ? "WebGPU display" : "Canvas fallback"} · wheel to zoom`; + cancelAnimationFrame(frame); + paint(); + updateColorControl(); + }) + .catch((error) => { + if (!disposed && current === generation) + status.textContent = `Display error: ${String(error)}`; + }); + } + function plotPosition(event: MouseEvent) { + const g = geometry(), + rect = canvas.getBoundingClientRect(); + const col = + (event.clientX - rect.left - g.left) / (g.width - g.left - g.right); + const row = + 1 - (event.clientY - rect.top - g.top) / (g.height - g.top - g.bottom); + return col >= 0 && col <= 1 && row >= 0 && row <= 1 ? [col, row] : null; + } + canvas.onpointerdown = (event) => { + if (event.button !== 0 || !plotPosition(event)) return; + event.preventDefault(); + event.stopPropagation(); + drag = { x: event.clientX, y: event.clientY, bounds: bounds.slice() }; + canvas.setPointerCapture(event.pointerId); + schedule(); + }; + canvas.onpointermove = (event) => { + const g = geometry(); + const width = g.width - g.left - g.right, + height = g.height - g.top - g.bottom; + if (drag) { + const dx = + ((event.clientX - drag.x) / width) * (drag.bounds[1] - drag.bounds[0]); + const dy = + ((event.clientY - drag.y) / height) * (drag.bounds[3] - drag.bounds[2]); + bounds = clampBounds([ + drag.bounds[0] - dx, + drag.bounds[1] - dx, + drag.bounds[2] + dy, + drag.bounds[3] + dy, + ]); + schedule(); + } + pointer = { clientX: event.clientX, clientY: event.clientY }; + showReadout(); + }; + function updateReadout() { + if (!pointer || !bitmap) { readout.title = readout.textContent = ""; return; } + const g = paintedGeometry, rect = canvas.getBoundingClientRect(); + const width = g.width - g.left - g.right, + height = g.height - g.top - g.bottom; + const colFraction = (pointer.clientX - rect.left - g.left) / width; + const rowFraction = 1 - (pointer.clientY - rect.top - g.top) / height; + if ( + colFraction < 0 || + colFraction >= 1 || + rowFraction < 0 || + rowFraction >= 1 + ) { + readout.title = readout.textContent = ""; + return; + } + const x = paintedBounds[0] + colFraction * (paintedBounds[1] - paintedBounds[0]), + y = paintedBounds[2] + rowFraction * (paintedBounds[3] - paintedBounds[2]); + const grid = model.get("grid"), + original = full(); + const col = Math.floor( + ((x - original[0]) / (original[1] - original[0])) * grid.cols, + ); + const row = Math.floor( + ((y - original[2]) / (original[3] - original[2])) * grid.rows, + ); + const xc = + original[0] + ((col + 0.5) * (original[1] - original[0])) / grid.cols; + const yc = + original[2] + ((row + 0.5) * (original[3] - original[2])) / grid.rows; + readout.title = readout.textContent = `Bin (${row}, ${col}) · x ${number(xc)} · y ${number(yc)} · value ${source[row * grid.cols + col]?.toPrecision(6)}`; + } + canvas.onpointerup = + canvas.onpointercancel = + canvas.onlostpointercapture = + () => { + if (!drag) return; + drag = null; + commit(); + schedule(); + }; + canvas.onpointerleave = () => { + pointer = null; + showReadout(); + }; + function zoomBy(factor: number, position = [0.5, 0.5]) { + bounds = clampBounds( + [0, 2].flatMap((i, axis) => { + const span = bounds[i + 1] - bounds[i]; + const anchor = bounds[i] + position[axis] * span; + return [ + anchor - position[axis] * span * factor, + anchor + (1 - position[axis]) * span * factor, + ]; + }), + ); + schedule(); + } + canvas.addEventListener( + "wheel", + (event) => { + const position = plotPosition(event); + if (!position) return; + event.preventDefault(); + event.stopPropagation(); + const delta = + event.deltaY * + (event.deltaMode === 1 + ? 16 + : event.deltaMode === 2 + ? geometry().height + : 1); + zoomBy(Math.exp(Math.max(-1, Math.min(1, delta * 0.002))), position); + clearTimeout(wheelTimer); + wheelTimer = window.setTimeout(commit, 150); + }, + { passive: false }, + ); + reset.onclick = () => { + clearTimeout(wheelTimer); + drag = null; + bounds = full().slice(); + commit(); + schedule(); + }; + canvas.ondblclick = reset.onclick as () => void; + zoomIn.onclick = () => { + zoomBy(1 / 1.5); + commit(); + }; + zoomOut.onclick = () => { + zoomBy(1.5); + commit(); + }; + save.onclick = () => + canvas.toBlob((blob) => { + if (blob) downloadBlob(blob, "plot2d.png"); + }); + const observers: [string, () => void][] = []; + for (const key of ["data_bytes", "cmap", "vmin", "vmax"]) + observers.push([`change:${key}`, prepare]); + for (const key of [ + "horizontal_line", + "title", + "x_label", + "y_label", + "colorbar_label", + "plot_height_px", + ]) + observers.push([`change:${key}`, schedule]); + observers.push(["change:max_width", size]); + observers.push(["change:cmap", updateColorControl]); + observers.push([ + "change:view_bounds", + () => { + bounds = ( + model.get("view_bounds").length ? model.get("view_bounds") : full() + ).slice(); + schedule(); + }, + ]); + observers.forEach(([event, handler]) => model.on(event, handler)); + const resize = new ResizeObserver(schedule); + resize.observe(host); + size(); + updateColorControl(); + prepare(); + return () => { + disposed = true; + generation++; + cancelAnimationFrame(frame); + clearTimeout(wheelTimer); + cancelAnimationFrame(readoutFrame); + colorRoot.unmount(); + resize.disconnect(); + observers.forEach(([event, handler]) => model.off(event, handler)); + if (bitmap instanceof ImageBitmap) bitmap.close(); + void Promise.all([queue, ready]).then(() => engine?.destroy()); + host.remove(); + }; +} + +export default { render }; diff --git a/js/plot2d/render.test.ts b/js/plot2d/render.test.ts new file mode 100644 index 00000000..ac6f56d0 --- /dev/null +++ b/js/plot2d/render.test.ts @@ -0,0 +1,193 @@ +import { act } from "react"; +import { afterEach, beforeEach, expect, it, vi } from "vitest"; +import plot2d from "./index"; + +const gpu = vi.hoisted(() => ({ + uploadData: vi.fn(), uploadLUT: vi.fn(), destroy: vi.fn(), create: vi.fn(), + renderSlotsToImageBitmapAsync: vi.fn(), +})); +vi.mock("../colormaps", () => ({ + COLORMAPS: { viridis: new Uint8Array(768), magma: new Uint8Array(768).fill(255) }, + createGPUColormapEngine: () => gpu.create(), + renderToOffscreen: vi.fn(() => document.createElement("canvas")), +})); + +class Bitmap { + close = vi.fn(); +} +class Model { + values: Record = { + grid: { rows: 2, cols: 3, bounds: [0, 3, 0, 2] }, + data_bytes: new Uint8Array(new Float64Array([1, 2, 3, 4, 5, 6]).buffer), + cmap: "viridis", vmin: 0, vmax: 6, view_bounds: [], + max_width: 600, plot_height_px: 350, + title: "Scalar map", x_label: "Distance", y_label: "Angle", colorbar_label: "Value", + }; + listeners = new Map void>>(); + get(key: string) { return this.values[key]; } + set(key: string, value: unknown) { + this.values[key] = value; + this.listeners.get(`change:${key}`)?.forEach(fn => fn()); + } + save_changes = vi.fn(); + on(key: string, fn: () => void) { + if (!this.listeners.has(key)) this.listeners.set(key, new Set()); + this.listeners.get(key)!.add(fn); + } + off(key: string, fn: () => void) { this.listeners.get(key)?.delete(fn); } +} +let cleanup: (() => void) | undefined; +let frames: Map; +let nextId: number; +let context: Record; +let el: HTMLDivElement; +let model: Model; +async function settle() { await act(async () => { for (let i = 0; i < 12; i++) await Promise.resolve(); }); } +function paintFrame() { + act(() => { const pending = [...frames.values()]; frames.clear(); pending.forEach(fn => fn(0)); }); +} +function deferred() { + let resolve!: (value: Bitmap[]) => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} +async function mount() { + act(() => { cleanup = plot2d.render({ model, el }); }); + await settle(); paintFrame(); +} +function hover() { + const canvas = el.querySelector("canvas")!; + canvas.onpointermove!({ clientX: 90, clientY: 190 } as PointerEvent); + paintFrame(); +} +beforeEach(() => { + vi.clearAllMocks(); + frames = new Map(); nextId = 0; + vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); + vi.stubGlobal("ImageBitmap", Bitmap); + vi.stubGlobal("ResizeObserver", class { observe() {} disconnect() {} }); + vi.stubGlobal("requestAnimationFrame", (fn: FrameRequestCallback) => { frames.set(++nextId, fn); return nextId; }); + vi.stubGlobal("cancelAnimationFrame", (id: number) => frames.delete(id)); + context = Object.fromEntries(["scale", "fillRect", "save", "translate", "rotate", "drawImage", "restore", "strokeRect", "fillText", "setLineDash", "beginPath", "moveTo", "lineTo", "stroke"].map(key => [key, vi.fn()])); + vi.spyOn(HTMLCanvasElement.prototype, "getContext").mockReturnValue(context as any); + document.body.dataset.jpThemeLight = "true"; + el = document.createElement("div"); document.body.append(el); model = new Model(); + gpu.create.mockResolvedValue(gpu); + gpu.renderSlotsToImageBitmapAsync.mockResolvedValue([new Bitmap()]); +}); +afterEach(async () => { + act(() => cleanup?.()); cleanup = undefined; + await settle(); el.remove(); delete document.body.dataset.jpThemeLight; + vi.restoreAllMocks(); vi.unstubAllGlobals(); +}); + +it("keeps visible pixels, hover and color limits together during rapid replacements", async () => { + await mount(); hover(); + expect(el.textContent).toContain("value 1.00000"); + const pending = deferred(), stale = new Bitmap(), latest = new Bitmap(); + gpu.renderSlotsToImageBitmapAsync.mockReturnValueOnce(pending.promise).mockResolvedValueOnce([latest]); + act(() => model.set("data_bytes", new Uint8Array(new Float64Array([11, 12, 13, 14, 15, 16]).buffer))); + await settle(); hover(); + expect(el.textContent).toContain("value 1.00000"); + act(() => { model.set("cmap", "magma"); model.set("vmax", 20); }); + paintFrame(); + expect(context.fillText).not.toHaveBeenCalledWith("20", expect.anything(), expect.anything()); + pending.resolve([stale]); await settle(); + expect(stale.close).toHaveBeenCalledOnce(); + expect(context.drawImage.mock.calls.some((call: unknown[]) => call[0] === stale)).toBe(false); + expect(context.drawImage.mock.lastCall?.[0]).toBe(latest); + expect(context.fillText).toHaveBeenCalledWith("20", expect.anything(), expect.anything()); + expect(el.textContent).toContain("value 11.0000"); + expect(gpu.uploadLUT.mock.lastCall?.[0]).toBe("magma"); +}); + +it("releases an in-flight bitmap and the engine when the view closes", async () => { + const pending = deferred(), bitmap = new Bitmap(); + gpu.renderSlotsToImageBitmapAsync.mockReturnValueOnce(pending.promise); + await mount(); + act(() => cleanup?.()); cleanup = undefined; + pending.resolve([bitmap]); await settle(); paintFrame(); + expect(bitmap.close).toHaveBeenCalledOnce(); + expect(gpu.destroy).toHaveBeenCalledOnce(); + expect(context.drawImage).not.toHaveBeenCalled(); + expect(el.childElementCount).toBe(0); + expect([...model.listeners.values()].every(set => set.size === 0)).toBe(true); +}); + +it("repaints notebook theme changes without changing scientific data", async () => { + await mount(); const bytes = model.get("data_bytes"); + const uploads = gpu.uploadData.mock.calls.length; + act(() => { document.body.dataset.jpThemeLight = "false"; }); + await settle(); paintFrame(); + expect(el.firstElementChild?.getAttribute("style")).toContain("rgb(30, 30, 30)"); + expect(context.fillStyle).toBe("#e0e0e0"); + expect(model.get("data_bytes")).toBe(bytes); + expect(gpu.uploadData).toHaveBeenCalledTimes(uploads); + expect(model.save_changes).not.toHaveBeenCalled(); +}); + +it("keeps zoom previews local and commits the final viewport on reset", async () => { + await mount(); + const bytes = model.get("data_bytes"), uploads = gpu.uploadData.mock.calls.length; + const canvas = el.querySelector("canvas")!; + canvas.dispatchEvent(new WheelEvent("wheel", { clientX: 150, clientY: 120, deltaY: -200, cancelable: true })); + paintFrame(); + expect(model.save_changes).not.toHaveBeenCalled(); + expect(model.get("view_bounds")).toEqual([]); + expect(context.drawImage.mock.lastCall?.[3]).toBeLessThan(3); + const reset = [...el.querySelectorAll("button")].find(button => button.textContent === "Reset View")!; + reset.click(); paintFrame(); + expect(model.get("view_bounds")).toEqual([0, 3, 0, 2]); + expect(context.drawImage.mock.lastCall?.slice(1, 5)).toEqual([0, 0, 3, 2]); + expect(model.get("data_bytes")).toBe(bytes); + expect(gpu.uploadData).toHaveBeenCalledTimes(uploads); +}); + +it("reports canvas fallback while retaining original readout values", async () => { + gpu.renderSlotsToImageBitmapAsync.mockResolvedValueOnce(null); + await mount(); hover(); + expect(el.textContent).toContain("Canvas fallback"); + expect(el.textContent).toContain("value 1.00000"); + expect(context.drawImage.mock.lastCall?.[0]).toBeInstanceOf(HTMLCanvasElement); +}); + +it("keeps a failed replacement from changing the displayed source", async () => { + await mount(); hover(); + gpu.renderSlotsToImageBitmapAsync.mockRejectedValueOnce(new Error("render interrupted")); + act(() => model.set("data_bytes", new Uint8Array(new Float64Array([21, 22, 23, 24, 25, 26]).buffer))); + await settle(); hover(); + expect(el.textContent).toContain("Display error: Error: render interrupted"); + expect(el.textContent).toContain("value 1.00000"); + act(() => model.set("cmap", "magma")); await settle(); + expect(el.textContent).toContain("value 21.0000"); + expect(el.textContent).not.toContain("Display error"); +}); + + +it("shows a lightweight saved preview until live data arrives", async () => { + model.values.data_bytes = new Uint8Array(); + model.values._static_fallback_jpeg = "iVBORw0KGgo="; + model.values._static_fallback_mime = "image/png"; + await mount(); + expect(el.querySelector("img")?.style.display).toBe("block"); + expect(el.querySelector("canvas")?.style.display).toBe("none"); + expect(gpu.uploadData).not.toHaveBeenCalled(); + expect(el.textContent).toContain("rerun the cell for interaction"); + act(() => model.set("data_bytes", new Uint8Array(new Float64Array([1, 2, 3, 4, 5, 6]).buffer))); + await settle(); hover(); + expect(el.querySelector("img")?.style.display).toBe("none"); + expect(el.querySelector("canvas")?.style.display).toBe("block"); + expect(el.textContent).toContain("value 1.00000"); +}); + + +it("releases a late engine when a saved-preview view closes before initialization", async () => { + let resolve!: (value: typeof gpu) => void; + gpu.create.mockReturnValueOnce(new Promise(done => { resolve = done; })); + model.values.data_bytes = new Uint8Array(); + await mount(); + act(() => cleanup?.()); cleanup = undefined; + expect(gpu.destroy).not.toHaveBeenCalled(); + resolve(gpu); await settle(); + expect(gpu.destroy).toHaveBeenCalledOnce(); +}); diff --git a/scripts/build.mjs b/scripts/build.mjs index dc3327aa..b82afd30 100644 --- a/scripts/build.mjs +++ b/scripts/build.mjs @@ -11,6 +11,7 @@ if (process.env.QUANTEM_WIDGET_SKIP_GPU_WEBGPU_SYNC !== "1") { syncGpuWebgpuSources(); } const widgets = [ + { name: "plot2d" }, { name: "show1d" }, { name: "show2d" }, { name: "mask2d" }, diff --git a/scripts/sync-gpu-webgpu.mjs b/scripts/sync-gpu-webgpu.mjs index d132ed37..4d03b26e 100644 --- a/scripts/sync-gpu-webgpu.mjs +++ b/scripts/sync-gpu-webgpu.mjs @@ -59,6 +59,21 @@ else: from importlib.resources import files root = files("quantem.gpu") +# New GPU versions retain display/webgpu/* as compatibility re-exports. +# Include the canonical owners when present; older wheels provide the complete +# implementations at the original paths. Keep the exported file list explicit. +for name in tuple(names): + if name.startswith("display/webgpu/"): + canonical = name.replace("display/webgpu/", "display/backends/webgpu/", 1) + if root.joinpath(*canonical.split("/")).is_file(): + names += (canonical,) + elif "/compute/webgpu/" in name: + canonical = name.replace("/compute/webgpu/", "/backends/webgpu/", 1) + if root.joinpath(*canonical.split("/")).is_file(): + names += (canonical,) +declarations = "io/backends/webgpu/jsfive.d.ts" +if root.joinpath(*declarations.split("/")).is_file(): + names += (declarations,) print(json.dumps({ name: root.joinpath(*name.split("/")).read_text(encoding="utf-8") for name in names diff --git a/scripts/widget_release_check.sh b/scripts/widget_release_check.sh index 10f3f6be..9af2808b 100755 --- a/scripts/widget_release_check.sh +++ b/scripts/widget_release_check.sh @@ -73,6 +73,7 @@ wheel = wheels[0] required = { "quantem/widget/static/chooselattice.js", "quantem/widget/static/show1d.js", + "quantem/widget/static/plot2d.js", "quantem/widget/static/show2d.js", "quantem/widget/static/show3d.js", "quantem/widget/static/show3dslices.js", diff --git a/src/quantem/widget/__init__.py b/src/quantem/widget/__init__.py index 0fea6e5d..050b1d7a 100644 --- a/src/quantem/widget/__init__.py +++ b/src/quantem/widget/__init__.py @@ -20,6 +20,7 @@ "ChooseLattice": ("quantem.widget.choose_lattice", "ChooseLattice"), "Mask2D": ("quantem.widget.mask2d", "Mask2D"), "Show1D": ("quantem.widget.show1d", "Show1D"), + "Plot2D": ("quantem.widget.plot2d", "Plot2D"), "Show2D": ("quantem.widget.show2d", "Show2D"), "Show3D": ("quantem.widget.show3d", "Show3D"), "Show3DSlices": ("quantem.widget.show3dslices", "Show3DSlices"), @@ -164,6 +165,7 @@ def free_gpu(verbose: bool = True) -> float: __all__ = [ "ChooseLattice", "Show1D", + "Plot2D", "Show2D", "Show3D", "Show3DSlices", diff --git a/src/quantem/widget/plot2d.py b/src/quantem/widget/plot2d.py new file mode 100644 index 00000000..34e24243 --- /dev/null +++ b/src/quantem/widget/plot2d.py @@ -0,0 +1,261 @@ +"""Calibrated scalar maps, distinct from spatial image viewers.""" + +import base64 +import io +from pathlib import Path +from typing import TYPE_CHECKING + +import anywidget +import numpy as np +import traitlets +from quantem.gpu.display import colormap_lut, colormap_names + +from .utils.array import _b64_safe +from .utils.static_fallback import StaticFallbackMixin + +if TYPE_CHECKING: + from matplotlib.figure import Figure + + +class Plot2D(StaticFallbackMixin, anywidget.AnyWidget): + """Inspect a scalar map with physical axes and a labeled color scale. + + Rows correspond to ``y`` and columns to ``x``. Both coordinates are bin + centers on increasing uniform grids. The first row is shown at the bottom, + as in Matplotlib's Cartesian plots. Source values and hover transport are + float64; QuantEM's shared WebGPU colormap renderer uses float32 display data. + Pan and zoom change only the viewport, never the scientific array. + + Parameters + ---------- + data : numpy.ndarray + Finite two-dimensional scalar values. + x, y : numpy.ndarray + Physical column and row bin centers, respectively. + x_label, y_label, colorbar_label : str + Scientific quantities including units where applicable. + title : str + Plot title. + cmap : str + QuantEM colormap name, also used for Matplotlib export. + vmin, vmax : float or None + Fixed color limits. When omitted, use the initial data range. + width, height : int + Initial plot dimensions in CSS pixels. + max_width : int + Maximum width in CSS pixels, also bounded by the notebook container. + + save_state : bool, default False + Embed the full float64 map in saved widget state when True. By default, + save a static PNG preview and omit the array from full state snapshots. + Targeted live updates always retain the original values. + + Examples + -------- + >>> plot = Plot2D(g3, x=radii, y=angles, x_label="r02 (Å)", + ... y_label="Angle (°)", colorbar_label="Normalized G3") + >>> plot.set_data(next_g3) + >>> plot.figure().savefig("g3.png", dpi=180) + """ + + _esm = Path(__file__).parent / "static" / "plot2d.js" + _save_state = traitlets.Bool(False).tag(sync=True) + _static_fallback_jpeg = traitlets.Unicode("").tag(sync=True) + _static_fallback_mime = traitlets.Unicode("image/png").tag(sync=True) + _UNSAVED_HEAVY_KEYS = ("data_bytes",) + data_bytes = traitlets.Bytes().tag(sync=True) + grid = traitlets.Dict().tag(sync=True) + title = traitlets.Unicode().tag(sync=True) + x_label = traitlets.Unicode().tag(sync=True) + y_label = traitlets.Unicode().tag(sync=True) + colorbar_label = traitlets.Unicode().tag(sync=True) + cmap = traitlets.Unicode("viridis").tag(sync=True) + vmin = traitlets.Float(0).tag(sync=True) + vmax = traitlets.Float(1).tag(sync=True) + horizontal_line = traitlets.Float(None, allow_none=True).tag(sync=True) + view_bounds = traitlets.List(traitlets.Float()).tag(sync=True) + plot_width_px = traitlets.Int(350).tag(sync=True) + plot_height_px = traitlets.Int(350).tag(sync=True) + max_width = traitlets.Int(600, min=260).tag(sync=True) + + def __init__( + self, + data: np.ndarray, + *, + x: np.ndarray, + y: np.ndarray, + x_label: str = "", + y_label: str = "", + colorbar_label: str = "Value", + title: str = "", + cmap: str = "viridis", + vmin: float | None = None, + vmax: float | None = None, + width: int = 350, + height: int = 350, + max_width: int = 600, + save_state: bool = False, + ) -> None: + self._save_state = bool(save_state) + self._configure_static_fallback(notebook_preview_format="png") + super().__init__() + self.x, self.y = ( + np.array(x, dtype=float, copy=True), + np.array(y, dtype=float, copy=True), + ) + edges = [] + for name, centers in (("x", self.x), ("y", self.y)): + if centers.ndim != 1 or len(centers) < 2 or not np.isfinite(centers).all(): + raise ValueError(f"{name} needs at least two finite bin centers.") + steps = np.diff(centers) + if not (steps > 0).all() or not np.allclose( + steps, steps[0], rtol=1e-5, atol=0.0 + ): + raise ValueError( + f"{name} must be an increasing uniform grid; got bin spacings " + f"from {steps.min():g} to {steps.max():g}. " + "Supply uniformly spaced bin centers in consistent units." + ) + edges.extend( + [float(centers[0] - steps[0] / 2), float(centers[-1] + steps[0] / 2)] + ) + self.grid = {"rows": len(self.y), "cols": len(self.x), "bounds": edges} + self.title, self.x_label, self.y_label = title, x_label, y_label + self.colorbar_label, self.cmap = colorbar_label, cmap + self.max_width = max_width + self.plot_width_px, self.plot_height_px = ( + min(max_width, max(260, width)), + max(260, height), + ) + self.layout.width = f"{self.plot_width_px}px" + self.layout.max_width = f"min(100%, {max_width}px)" + self.set_data(data) + self.vmin = float(self.data.min()) if vmin is None else float(vmin) + self.vmax = float(self.data.max()) if vmax is None else float(vmax) + if self.vmax == self.vmin and vmin is None and vmax is None: + self.vmax = self.vmin + 1 + if not np.isfinite([self.vmin, self.vmax]).all() or self.vmax <= self.vmin: + raise ValueError("Use finite color limits with vmax > vmin.") + + def get_state(self, key=None, drop_defaults=False): + """Return full saved state or an untrimmed targeted live update.""" + state = super().get_state(key=key, drop_defaults=drop_defaults) + if key is None and not self._save_state: + state.pop("data_bytes", None) + if self._static_fallback_enabled(): + preview = self._static_fallback_png_b64() + if preview: + self._store_static_fallback_preview(preview) + state["_static_fallback_jpeg"] = self._static_fallback_jpeg + state["_static_fallback_mime"] = self._static_fallback_mime + else: + state.pop("_static_fallback_jpeg", None) + state.pop("_static_fallback_mime", None) + return state + + def _store_static_fallback_preview(self, png_b64: str) -> None: + """Retain a lossless preview for lightweight model restoration.""" + if not self._save_state: + self._static_fallback_jpeg = png_b64 + self._static_fallback_mime = "image/png" + + def _static_png_b64(self, max_px: int = 512) -> str | None: + """Render a bounded saved preview without modifying scientific data.""" + if not hasattr(self, "data"): + return None + figure = self._make_figure(max_bins=max_px) + buffer = io.BytesIO() + figure.savefig(buffer, format="png", dpi=100) + return base64.b64encode(buffer.getvalue()).decode("ascii") + + @traitlets.validate("cmap") + def _validate_cmap(self, proposal: dict) -> str: + value = proposal["value"] + if value not in colormap_names(): + raise ValueError( + f"Unknown colormap {value!r}; choose from {colormap_names()}." + ) + return value + + def set_data(self, data: np.ndarray) -> None: + """Replace values on the existing grid, preserving limits and viewport. + + Parameters + ---------- + data : numpy.ndarray + Finite values with the original row/column shape. + + Examples + -------- + >>> plot.set_data(predicted_g3[1]) + """ + if np.iscomplexobj(data): + raise ValueError( + "Choose a real scalar quantity explicitly before plotting." + ) + values = np.asarray(data, dtype=np.float64) + if values.shape != (len(self.y), len(self.x)) or not np.isfinite(values).all(): + raise ValueError( + f"Supply finite data shaped {(len(self.y), len(self.x))}; got {values.shape}." + ) + self.data = values.copy() + self.data_bytes = _b64_safe(np.ascontiguousarray(self.data).tobytes()) + + def figure(self) -> "Figure": + """Return a closed Matplotlib figure of the current values and viewport. + + Returns + ------- + matplotlib.figure.Figure + Editable figure with physical axes and a labeled colorbar. + + Examples + -------- + >>> plot.figure().savefig("correlation.svg") + """ + return self._make_figure() + + def _make_figure(self, max_bins: int | None = None) -> "Figure": + """Draw exact data for export or a strided, bounded notebook preview.""" + import matplotlib.pyplot as plt + from matplotlib.colors import ListedColormap + + size = (5, 4) if max_bins is None else (max_bins / 100, max_bins * 0.8 / 100) + figure, axes = plt.subplots(figsize=size, layout="constrained") + values = self.data + x, y = self.x, self.y + shading = "nearest" + if max_bins is not None: + rows, cols = values.shape + row_step = max(1, int(np.ceil(rows / max_bins))) + col_step = max(1, int(np.ceil(cols / max_bins))) + # Explicit edges preserve calibration, including the last partial bin. + col_edges = np.r_[np.arange(0, cols, col_step), cols] + row_edges = np.r_[np.arange(0, rows, row_step), rows] + left, right, bottom, top = self.grid["bounds"] + x = left + col_edges * (right - left) / cols + y = bottom + row_edges * (top - bottom) / rows + values = values[::row_step, ::col_step] + shading = "flat" + mesh = axes.pcolormesh( + x, + y, + values, + shading=shading, + cmap=ListedColormap(colormap_lut(self.cmap), name=self.cmap), + vmin=self.vmin, + vmax=self.vmax, + ) + limits = self.view_bounds or self.grid["bounds"] + axes.set( + xlabel=self.x_label, + ylabel=self.y_label, + title=self.title, + xlim=limits[:2], + ylim=limits[2:], + ) + if self.horizontal_line is not None: + axes.axhline(self.horizontal_line, color="#e649a0", linestyle=":") + figure.colorbar(mesh, ax=axes, label=self.colorbar_label) + plt.close(figure) + return figure diff --git a/tests/plot2d/test_plot2d.py b/tests/plot2d/test_plot2d.py new file mode 100644 index 00000000..f0016480 --- /dev/null +++ b/tests/plot2d/test_plot2d.py @@ -0,0 +1,122 @@ +"""Review a calibrated correlation map, change its window, and export it.""" + +import base64 + +import matplotlib.pyplot as plt +import numpy as np +import pytest +from ipywidgets import Widget +from ipywidgets.embed import embed_data + +from quantem.widget import Plot2D + + +def test_calibrated_map_export_and_saved_state(tmp_path): + radius = (np.arange(100) + 0.5) * 0.1 + angles = (np.arange(36) + 0.5) * 5 + truth = np.cos(np.deg2rad(angles[:, None])) ** 2 * radius[None, :] + plot = Plot2D( + truth, + x=radius, + y=angles, + x_label="r02 (Å)", + y_label="Angle (°)", + colorbar_label="Normalized G3", + width=1200, + max_width=420, + save_state=True, + ) + assert plot.plot_width_px == 420 + assert plot.layout.max_width == "min(100%, 420px)" + original = truth.copy() + np.testing.assert_allclose(plot.grid["bounds"], [0, 10, 0, 180], atol=1e-14) + plot.view_bounds = [1, 8, 30, 150] + plot.horizontal_line = 92.5 + plot.set_data(truth * 0.75) + fixed_limits = (plot.vmin, plot.vmax) + plot.cmap = "magma" + figures_before = plt.get_fignums() + figure = plot.figure() + assert plt.get_fignums() == figures_before + axes = figure.axes[0] + np.testing.assert_array_equal(axes.collections[0].get_array(), truth * 0.75) + np.testing.assert_array_equal(axes.get_xlim(), [1, 8]) + np.testing.assert_array_equal(axes.get_ylim(), [30, 150]) + assert axes.get_xlabel() == "r02 (Å)" + assert figure.axes[1].get_ylabel() == "Normalized G3" + assert axes.collections[0].cmap.name == "magma" + assert axes.collections[0].get_clim() == fixed_limits + figure.savefig(tmp_path / "g3.svg") + own_state = Widget.get_manager_state(widgets=[plot, plot.layout])["state"] + state = embed_data(views=[plot], state=own_state)["manager_state"]["state"][ + plot.model_id + ] + assert state["state"]["horizontal_line"] == 92.5 + buffer = next(item for item in state["buffers"] if item["path"] == ["data_bytes"]) + restored = np.frombuffer( + base64.b64decode(buffer["data"]), dtype=np.float64, count=truth.size + ) + np.testing.assert_array_equal(restored, (truth * 0.75).ravel()) + np.testing.assert_array_equal(truth, original) + + +def test_asymmetric_map_retains_row_column_calibration(): + """Export a nonsquare, strided map without swapping axes or changing values.""" + radius = np.array([1.0, 1.5, 2.0]) + angle = np.array([30.0, 90.0]) + values = np.arange(12, dtype=np.float64).reshape(2, 6)[:, ::2] + plot = Plot2D(values, x=radius, y=angle, vmin=-1, vmax=12) + figure = plot.figure() + mesh = figure.axes[0].collections[0] + expected_column_edges = [0.75, 1.25, 1.75, 2.25] + expected_row_edges = [0.0, 60.0, 120.0] + np.testing.assert_array_equal( + mesh.get_coordinates()[0, :, 0], expected_column_edges + ) + np.testing.assert_array_equal(mesh.get_coordinates()[:, 0, 1], expected_row_edges) + np.testing.assert_array_equal(mesh.get_array(), values) + # Reusing caller-owned input arrays cannot change the already-created map. + values[:] = -1 + radius[:] = 0 + angle[:] = 0 + np.testing.assert_array_equal( + plot.figure().axes[0].collections[0].get_array(), [[0, 2, 4], [6, 8, 10]] + ) + + +def test_meter_calibration_preserves_uniform_grid_requirement(): + """Changing distance units to meters must not allow misleading axes.""" + distances_m = np.array([1, 2, 3]) * 1e-10 + angles = np.array([30, 90]) + values = np.arange(6).reshape(2, 3) + plot = Plot2D(values, x=distances_m, y=angles, x_label="Distance (m)") + np.testing.assert_allclose( + plot.grid["bounds"][:2], [0.5e-10, 3.5e-10], rtol=1e-14, atol=0 + ) + for invalid in ([0, 1e-10, 5e-11], [0, 1e-10, 1e-10], [0, 1e-10, 5e-9]): + with pytest.raises(ValueError, match="uniformly spaced bin centers"): + Plot2D(values, x=np.array(invalid), y=angles) + + +def test_lightweight_snapshot_preserves_live_values_and_current_preview(): + """Save a calibrated view without embedding its source array by default.""" + values = np.arange(24, dtype=float).reshape(4, 6) + plot = Plot2D(values, x=np.arange(6), y=np.arange(4)) + first = plot.get_state() + assert "data_bytes" not in first + assert first["_static_fallback_mime"] == "image/png" + assert base64.b64decode(first["_static_fallback_jpeg"]).startswith(b"\x89PNG") + plot.view_bounds = [1, 4, 0, 2] + plot.horizontal_line = 1.5 + plot.set_data(values[::-1]) + saved = plot.get_state() + assert saved["_static_fallback_jpeg"] != first["_static_fallback_jpeg"] + assert saved["view_bounds"] == [1, 4, 0, 2] + assert saved["horizontal_line"] == 1.5 + assert "data_bytes" not in saved + live = plot.get_state({"data_bytes", "grid"}) + np.testing.assert_array_equal( + np.frombuffer(live["data_bytes"], dtype=np.float64, count=values.size), + values[::-1].ravel(), + ) + plot.close() diff --git a/tests/test_save_state.py b/tests/test_save_state.py index edee7eb5..c4a925f8 100644 --- a/tests/test_save_state.py +++ b/tests/test_save_state.py @@ -22,7 +22,7 @@ import pytest from PIL import Image -from quantem.widget import Show1D, Show2D, Show3D, Show4DSTEM, ShowEDS +from quantem.widget import Plot2D, Show1D, Show2D, Show3D, Show4DSTEM, ShowEDS IMAGE_MIME_KEYS = ("image/jpeg", "image/webp", "image/png") @@ -70,6 +70,10 @@ def _mos2_like_stack(frames: int, rows: int, cols: int) -> np.ndarray: def _make(widget, *, save_state): """Construct a small instance of each widget plus the trait key that carries its live-render pixels (the one that must survive the targeted send path).""" + if widget is Plot2D: + data = np.arange(24, dtype=float).reshape(4, 6) + return Plot2D(data, x=np.arange(6), y=np.arange(4), + save_state=save_state), "data_bytes" if widget is Show1D: # snapshot_bytes is empty on a plain trace widget but the KEY must # still survive the targeted path (it streams when a monitor attaches). @@ -88,7 +92,7 @@ def _make(widget, *, save_state): save_state=save_state), "virtual_image_bytes" -WIDGETS = [Show1D, Show2D, Show3D, Show4DSTEM, ShowEDS] +WIDGETS = [Plot2D, Show1D, Show2D, Show3D, Show4DSTEM, ShowEDS] @pytest.mark.parametrize("widget", WIDGETS) @@ -101,7 +105,7 @@ def test_targeted_send_state_never_trimmed(widget): assert render_key in w.get_state(render_key), ( f"{widget.__name__}: targeted get_state({render_key!r}) dropped the key " f"- live render would go blank") - assert render_key in w.get_state({render_key, "widget_version"}), ( + assert render_key in w.get_state({render_key, "layout"}), ( f"{widget.__name__}: hold_sync batch lost {render_key!r}") @@ -127,8 +131,8 @@ def test_static_fallback_present(widget): bundle = w._repr_mimebundle_() data = bundle[0] if isinstance(bundle, tuple) else bundle image_keys = [key for key in IMAGE_MIME_KEYS if key in (data or {})] - assert image_keys == ["image/jpeg"], ( - f"{widget.__name__}: no default JPEG static fallback for a cold reopen") + assert image_keys == ["image/png" if widget is Plot2D else "image/jpeg"], ( + f"{widget.__name__}: missing expected static fallback for a cold reopen") @pytest.mark.parametrize("widget", WIDGETS) @@ -1752,7 +1756,9 @@ def fake_display(data, raw=False, metadata=None, display_id=None, **kw): assert png_calls, f"{widget.__name__}: deferred fill never rendered the PNG" assert "post_execute" not in hooks, "one-shot hook did not unregister" fill_data, fill_meta = updated[-1] - assert isinstance(fill_data["image/jpeg"], bytes) and len(fill_data["image/jpeg"]) > 1000 + preview_mime = "image/png" if widget is Plot2D else "image/jpeg" + assert isinstance(fill_data[preview_mime], bytes) + assert len(fill_data[preview_mime]) > 1000 assert "quantem-static-fallback" in fill_data["text/html"] assert fill_meta == {"quantem.widget": {"static_fallback": True}} # the saved-notebook snapshot must stay free of the bulk buffers