diff --git a/js/scanRotationParity.test.ts b/js/scanRotationParity.test.ts new file mode 100644 index 00000000..012e6c85 --- /dev/null +++ b/js/scanRotationParity.test.ts @@ -0,0 +1,56 @@ +import { describe, expect, it } from "vitest"; +import fixture from "./.generated/engine/parity/scan_rotation_v1.json"; +import { + SCAN_QUARTER_TURN_WGSL, + type ScanQuarterTurns, + scanQuarterTurnOutputShape, + scanQuarterTurnSourceIndex, +} from "./.generated/engine/geometry/compute/webgpu/quarter-turn"; + +describe("shared scan-rotation gold fixture", () => { + it("maps every detector pattern with the canonical row-column convention", () => { + const [scanRows, scanColumns, detectorRows, detectorColumns] = fixture.source.shape; + const detectorPixels = detectorRows * detectorColumns; + + for (const testCase of fixture.cases) { + const quarterTurns = testCase.quarter_turns_counterclockwise as ScanQuarterTurns; + const [outputRows, outputColumns] = scanQuarterTurnOutputShape( + scanRows, + scanColumns, + quarterTurns, + ); + expect([outputRows, outputColumns, detectorRows, detectorColumns]).toEqual( + testCase.output_shape, + ); + for (let outputRow = 0; outputRow < outputRows; outputRow++) { + for (let outputColumn = 0; outputColumn < outputColumns; outputColumn++) { + const sourceScan = scanQuarterTurnSourceIndex( + outputRow, + outputColumn, + scanRows, + scanColumns, + quarterTurns, + ); + const outputScan = outputRow * outputColumns + outputColumn; + const sourceStart = sourceScan * detectorPixels; + const outputStart = outputScan * detectorPixels; + expect( + testCase.expected_values.slice(outputStart, outputStart + detectorPixels), + ).toEqual( + fixture.source.values.slice(sourceStart, sourceStart + detectorPixels), + ); + } + } + } + }); + + it("keeps the same mapping in the hardware shader contract", () => { + expect(SCAN_QUARTER_TURN_WGSL).toContain("sourceRow = outputColumn"); + expect(SCAN_QUARTER_TURN_WGSL).toContain( + "sourceRow = parameters.sourceRows - 1u - outputColumn", + ); + expect(SCAN_QUARTER_TURN_WGSL).toContain( + "sourceScan * parameters.wordsPerScan + wordInScan", + ); + }); +}); diff --git a/js/show4dstem/detectorInteraction.test.ts b/js/show4dstem/detectorInteraction.test.ts new file mode 100644 index 00000000..d0fd99c7 --- /dev/null +++ b/js/show4dstem/detectorInteraction.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it } from "vitest"; +import { clampDetectorCenter, resizeDetectorFromPointer } from "./detectorInteraction"; + +describe("Show4DSTEM detector interaction geometry", () => { + it("keeps subpixel detector centers while clamping to the diffraction plane", () => { + expect(clampDetectorCenter(12.25, 18.75, 48, 48)).toEqual({ + row: 12.25, + col: 18.75, + }); + expect(clampDetectorCenter(-2, 50, 48, 48)).toEqual({ row: 0, col: 47 }); + }); + + it("resizes circle, square, and rectangle detectors from the live pointer", () => { + const common = { + centerRow: 10, + centerCol: 10, + pointerRow: 13, + pointerCol: 14, + radius: 8, + radiusInner: 3, + }; + + expect(resizeDetectorFromPointer({ ...common, mode: "circle" })).toEqual({ radius: 5 }); + expect(resizeDetectorFromPointer({ ...common, mode: "square" })).toEqual({ radius: 4 }); + expect(resizeDetectorFromPointer({ ...common, mode: "rect" })).toEqual({ + width: 8, + height: 6, + }); + expect(resizeDetectorFromPointer({ + ...common, + mode: "rect", + aspectRatio: 2, + preserveAspect: true, + })).toEqual({ width: 12, height: 6 }); + }); + + it("keeps annular inner and outer radii ordered during live resizing", () => { + const common = { + mode: "annular" as const, + centerRow: 10, + centerCol: 10, + pointerRow: 10, + pointerCol: 12, + radius: 10, + radiusInner: 4, + }; + + expect(resizeDetectorFromPointer(common)).toEqual({ radius: 5 }); + expect(resizeDetectorFromPointer({ + ...common, + pointerCol: 25, + resizeInner: true, + })).toEqual({ radiusInner: 9 }); + }); +}); diff --git a/js/show4dstem/detectorInteraction.ts b/js/show4dstem/detectorInteraction.ts new file mode 100644 index 00000000..af93e393 --- /dev/null +++ b/js/show4dstem/detectorInteraction.ts @@ -0,0 +1,74 @@ +export type DetectorRoiMode = "point" | "circle" | "square" | "rect" | "annular" | "off"; + +export type DetectorResizeGeometry = { + radius?: number; + radiusInner?: number; + width?: number; + height?: number; +}; + +export function clampDetectorCenter( + row: number, + col: number, + detectorRows: number, + detectorCols: number, +): { row: number; col: number } { + return { + row: Math.max(0, Math.min(detectorRows - 1, row)), + col: Math.max(0, Math.min(detectorCols - 1, col)), + }; +} + +export function resizeDetectorFromPointer({ + mode, + centerRow, + centerCol, + pointerRow, + pointerCol, + radius, + radiusInner, + resizeInner = false, + aspectRatio = null, + preserveAspect = false, +}: { + mode: DetectorRoiMode; + centerRow: number; + centerCol: number; + pointerRow: number; + pointerCol: number; + radius: number; + radiusInner: number; + resizeInner?: boolean; + aspectRatio?: number | null; + preserveAspect?: boolean; +}): DetectorResizeGeometry | null { + const rowDistance = Math.abs(pointerRow - centerRow); + const colDistance = Math.abs(pointerCol - centerCol); + + if (resizeInner && mode === "annular") { + return { + radiusInner: Math.max(1, Math.min(radius - 1, Math.hypot(rowDistance, colDistance))), + }; + } + + if (mode === "rect") { + let width = Math.max(2, colDistance * 2); + let height = Math.max(2, rowDistance * 2); + if (preserveAspect && aspectRatio != null) { + if (width / height > aspectRatio) height = Math.max(2, width / aspectRatio); + else width = Math.max(2, height * aspectRatio); + } + return { width, height }; + } + + if (mode === "circle" || mode === "square" || mode === "annular") { + const nextRadius = mode === "square" + ? Math.max(rowDistance, colDistance) + : Math.hypot(rowDistance, colDistance); + return { + radius: Math.max(mode === "annular" ? radiusInner + 1 : 1, nextRadius), + }; + } + + return null; +} diff --git a/js/show4dstem/index.tsx b/js/show4dstem/index.tsx index 4263fa62..f56a4d58 100644 --- a/js/show4dstem/index.tsx +++ b/js/show4dstem/index.tsx @@ -69,6 +69,11 @@ import { type ComparePageMessage, type ProgressiveComparePage, } from "./progressiveCompare"; +import { + clampDetectorCenter, + resizeDetectorFromPointer, + type DetectorRoiMode, +} from "./detectorInteraction"; function normaliseViSource(value: unknown): string { const raw = String(value || "roi").trim(); @@ -1491,6 +1496,7 @@ interface CompareVirtualGridProps { // GPU-resident panels: frame -> engine colormap slot, painted with a GPU range // through each tile's visible WebGPU canvas; bytes stay the settle/export fallback. gpuSlots?: Map | null; + gpuRanges?: Map | null; gpuVersion?: number; gpuEngine?: GPUColormapEngine | null; progressivePage?: ProgressiveComparePage | null; @@ -1542,6 +1548,7 @@ function CompareVirtualGrid({ count, indices, gpuSlots, + gpuRanges, gpuVersion, gpuEngine, progressivePage, @@ -1586,7 +1593,7 @@ function CompareVirtualGrid({ }: CompareVirtualGridProps) { const canvasRefs = React.useRef<(HTMLCanvasElement | null)[]>([]); const gpuCanvasRefs = React.useRef<(HTMLCanvasElement | null)[]>([]); - const gpuCanvasContextsRef = React.useRef<(GPUCanvasContext | null)[]>([]); + const gpuRenderGenerationRef = React.useRef(0); const canvasDrawCacheRef = React.useRef(new Map ({ frame, panel: panelByFrame.get(frame), - gpuLoaded: Boolean(gpuSlots?.has(frame) && gpuEngine), + gpuLoaded: Boolean(gpuSlots?.has(frame) && gpuRanges?.has(frame) && gpuEngine), })) .filter((entry) => Boolean(progressivePage) || entry.panel !== undefined || entry.gpuLoaded); - }, [gpuEngine, gpuSlots, gpuVersion, panelByFrame, progressivePage, renderIndices, scaleMode]); + }, [gpuEngine, gpuRanges, gpuSlots, gpuVersion, panelByFrame, progressivePage, renderIndices, scaleMode]); const renderGpuSlotsNow = React.useCallback((): number => { if (!gpuEngine || !gpuSlots) return 0; const lut = COLORMAPS[colormap] || COLORMAPS.inferno; gpuEngine.uploadLUT(colormap, lut); - let painted = 0; + const generation = ++gpuRenderGenerationRef.current; + const panels: { + canvas: HTMLCanvasElement; + range: { vmin: number; vmax: number }; + slot: number; + }[] = []; renderEntries.forEach((entry, localIdx) => { const slot = gpuSlots.get(entry.frame); + const rawRange = gpuRanges?.get(entry.frame); const canvas = gpuCanvasRefs.current[localIdx]; - if (slot === undefined || !canvas) return; - if ( - canvas.width !== shapeCols - || canvas.height !== shapeRows - || !gpuCanvasContextsRef.current[localIdx] - ) { - gpuCanvasContextsRef.current[localIdx] = gpuEngine.configureCanvas(canvas, shapeCols, shapeRows); - } - const ctx = gpuCanvasContextsRef.current[localIdx]; - if (!ctx) return; - const ok = gpuEngine.renderSlotDirectWithGpuRangeToCanvas( - slot, - vminPct, - vmaxPct, - scaleMode === "log", - ctx, + if (slot === undefined || !rawRange || !canvas) return; + const transformRangeValue = (value: number) => scaleMode === "log" + ? (value >= 0 ? Math.log1p(value) : -Math.log1p(-value)) + : value; + const rangeMin = transformRangeValue(rawRange.min); + const rangeMax = transformRangeValue(rawRange.max); + const span = Math.max(0, rangeMax - rangeMin); + const displayRange = { + vmin: rangeMin + span * Math.max(0, Math.min(100, vminPct)) / 100, + vmax: rangeMin + span * Math.max(0, Math.min(100, vmaxPct)) / 100, + }; + panels.push({ canvas, range: displayRange, slot }); + }); + if (!panels.length) return 0; + void (async () => { + const bitmap = await gpuEngine.renderPanelSlotsToImageBitmapAsync( + panels.map((panel) => panel.slot), + panels.map((panel) => panel.range), + panels.map(() => scaleMode === "log"), { - width: shapeCols, + width: shapeCols * panels.length, height: shapeRows, + panelCount: panels.length, + cols: panels.length, + rows: 1, + gap: 0, bgRgb: 0, - transform: { zoom: compareZoom, panX: comparePanX, panY: comparePanY }, + transforms: panels.map(() => ({ + zoom: compareZoom, + panX: comparePanX, + panY: comparePanY, + })), smooth, }, ); - if (ok) painted++; - }); - if (painted > 0) onGpuPaint?.(painted); - return painted; - }, [colormap, comparePanX, comparePanY, compareZoom, gpuEngine, gpuSlots, onGpuPaint, renderEntries, scaleMode, shapeCols, shapeRows, smooth, vmaxPct, vminPct]); + if (!bitmap || generation !== gpuRenderGenerationRef.current) { + bitmap?.close(); + return; + } + let painted = 0; + panels.forEach((panel, index) => { + if (!panel.canvas.isConnected) return; + if (panel.canvas.width !== shapeCols) panel.canvas.width = shapeCols; + if (panel.canvas.height !== shapeRows) panel.canvas.height = shapeRows; + const context = panel.canvas.getContext("2d"); + if (!context) return; + context.imageSmoothingEnabled = false; + context.clearRect(0, 0, shapeCols, shapeRows); + context.drawImage( + bitmap, + index * shapeCols, + 0, + shapeCols, + shapeRows, + 0, + 0, + shapeCols, + shapeRows, + ); + painted++; + }); + bitmap.close(); + if (painted > 0) onGpuPaint?.(painted); + })(); + return panels.length; + }, [colormap, comparePanX, comparePanY, compareZoom, gpuEngine, gpuRanges, gpuSlots, onGpuPaint, renderEntries, scaleMode, shapeCols, shapeRows, smooth, vmaxPct, vminPct]); React.useEffect(() => { onGpuRendererReady?.(renderGpuSlotsNow); @@ -2312,7 +2362,6 @@ function CompareVirtualGrid({ { gpuCanvasRefs.current[localIdx] = node; - if (!node) gpuCanvasContextsRef.current[localIdx] = null; }} width={shapeCols} height={shapeRows} @@ -3131,10 +3180,12 @@ function Show4DSTEM() { const [viGpuVersion, setViGpuVersion] = React.useState(0); const [viGpuRetainedReady, setViGpuRetainedReady] = React.useState(false); const viGpuImageRef = React.useRef(null); - // GPU-resident compare panels: frame index -> engine colormap slot. Written by - // the interactive compare recompute (no readback), consumed by the grid painter. + // GPU-resident compare panels: frame index -> engine colormap slot. Only a + // small settled min/max reduction is read back; scientific image pixels stay + // on the GPU and the interactive drag path reuses the cached range. const [compareGpuVersion, setCompareGpuVersion] = React.useState(0); const compareGpuSlotsRef = React.useRef(new Map()); + const compareGpuRangesRef = React.useRef(new Map()); const compareGpuHistogramGenRef = React.useRef(0); const compareGpuRenderNowRef = React.useRef<(() => number) | null>(null); const compareIncrementalRef = React.useRef<{ @@ -3150,10 +3201,11 @@ function Show4DSTEM() { lastAdoptedPanels: 0, lastRequestedPanels: 0, lastPaintedPanels: 0, + lastRangeReadbackBytes: 0, }); const publishLiveCompareViStats = React.useCallback(( event: string, - detail: { ms?: number; adoptedPanels?: number; requestedPanels?: number; paintedPanels?: number; addedPixels?: number; removedPixels?: number }, + detail: { ms?: number; adoptedPanels?: number; requestedPanels?: number; paintedPanels?: number; addedPixels?: number; removedPixels?: number; rangeReadbackBytes?: number }, ) => { const now = performance.now(); const stats = liveCompareViStatsRef.current; @@ -3162,6 +3214,7 @@ function Show4DSTEM() { stats.lastComputeMs = detail.ms ?? 0; stats.lastAdoptedPanels = detail.adoptedPanels ?? 0; stats.lastRequestedPanels = detail.requestedPanels ?? 0; + stats.lastRangeReadbackBytes = detail.rangeReadbackBytes ?? 0; } else { stats.paintTimes.push(now); stats.lastPaintMs = now; @@ -3174,7 +3227,8 @@ function Show4DSTEM() { const recentPaint = stats.paintTimes.length; const payload = { event, - gpuOnlyHotPath: true, + gpuOnlyHotPath: stats.lastRangeReadbackBytes === 0, + rangeReadbackBytes: stats.lastRangeReadbackBytes, computeFps: Math.round(recentCompute * 10) / 10, paintFps: Math.round(recentPaint * 10) / 10, lastComputeMs: Math.round(stats.lastComputeMs * 10) / 10, @@ -4673,6 +4727,7 @@ function Show4DSTEM() { compareIncrementalRef.current = null; if (compareGpuSlotsRef.current.size) { compareGpuSlotsRef.current.clear(); + compareGpuRangesRef.current.clear(); setCompareGpuVersion(v => v + 1); } }; @@ -4721,6 +4776,7 @@ function Show4DSTEM() { compareGpuSlotsRef.current.set(idx, slot); } let adopted = 0; + let rangeReadbackBytes = 0; if (batchComputes.length) { const indicesKey = batchFrames.join(","); const previous = compareIncrementalRef.current; @@ -4768,6 +4824,15 @@ function Show4DSTEM() { nextBuffers.set(batchFrames[i], buffers[i]); adopted++; } + const rangesReady = batchFrames.every((frame) => compareGpuRangesRef.current.has(frame)); + if (!interactiveDrag || !rangesReady) { + const ranges = await engine0.computeRangeBatch(batchSlots); + rangeReadbackBytes = batchSlots.length * Math.ceil((scanRows * scanCols) / 256) * 2 * 4; + ranges.forEach((range, index) => { + const frame = batchFrames[index]; + if (frame !== undefined) compareGpuRangesRef.current.set(frame, range); + }); + } compareIncrementalRef.current = { mask: new Uint32Array(mask0), buffers: nextBuffers, @@ -4784,6 +4849,7 @@ function Show4DSTEM() { addedPixels, removedPixels, paintedPanels: paintedNow, + rangeReadbackBytes, }); } if (adopted) bumpCompareGpuVersion(); @@ -4859,6 +4925,14 @@ function Show4DSTEM() { // model.set a silent no-op (no change event -> stats/export/save-state stale) publishDirectCompareStack(new DataView(stack.slice().buffer), indices.length, indices); }; + const recomputeVisibleVirtualImages = async () => { + const mode = String(model.get("view_mode") || "single"); + if (mode === "multiple" || mode === "compare") { + await recomputeCompareVI(); + return; + } + await recomputeVI(); + }; (window as unknown as { __sh4d: unknown }).__sh4d = { model, recomputeVI, recomputeCompareVI, detMask: () => buildDetectorMask(model, detR, detC), deriveOnly: async () => { const vi = await compute!.maskedSum(buildDetectorMask(model, detR, detC)); return vi.length; }, @@ -5445,10 +5519,7 @@ function Show4DSTEM() { flushRoiCenter(); flushRoiRadius(); compareViLiveInFlightRef.current = true; - void (async () => { - await recomputeVI(); - await recomputeCompareVI(); - })().finally(() => { + void recomputeVisibleVirtualImages().finally(() => { compareViLiveInFlightRef.current = false; if (compareViLivePendingRef.current && dpRoiInteractiveRef.current) { compareViLivePendingRef.current = false; @@ -5459,8 +5530,7 @@ function Show4DSTEM() { }); }; requestViFinalizeRef.current = () => { - void recomputeVI(); - void recomputeCompareVI(); + void recomputeVisibleVirtualImages(); }; const recomputeDP = async () => { const mode = model.get("vi_roi_mode"); @@ -5527,9 +5597,15 @@ function Show4DSTEM() { }; if (h5VolumePreload) { void h5VolumePreload.then(() => { - if (disposed) return; - void recomputeCompareVI(); - void recomputeFrame(); + const refreshLoadedH5Views = () => { + if (disposed) return; + void recomputeCompareVI(); + void recomputeFrame(); + }; + refreshLoadedH5Views(); + requestAnimationFrame(() => { + requestAnimationFrame(refreshLoadedH5Views); + }); }).catch((error) => { console.warn("Show4DSTEM HDF5 volume preload refresh failed", error); }); @@ -8406,31 +8482,40 @@ function Show4DSTEM() { const resizeDpRoiFromImagePoint = React.useCallback((imgX: number, imgY: number, shiftKey: boolean = false): boolean => { if (isDraggingResizeInner) { - const dx = Math.abs(imgX - activeRoiCenterCol); - const dy = Math.abs(imgY - activeRoiCenterRow); - const newRadius = Math.sqrt(dx ** 2 + dy ** 2); - setRoiRadiusInner(Math.max(1, Math.min(roiRadius - 1, Math.round(newRadius)))); + const geometry = resizeDetectorFromPointer({ + mode: roiMode as DetectorRoiMode, + centerRow: activeRoiCenterRow, + centerCol: activeRoiCenterCol, + pointerRow: imgY, + pointerCol: imgX, + radius: roiRadius, + radiusInner: roiRadiusInner, + resizeInner: true, + }); + if (geometry?.radiusInner === undefined) return false; + setRoiRadiusInner(geometry.radiusInner); requestCompareViLive(); return true; } if (isDraggingResize) { - const dx = Math.abs(imgX - activeRoiCenterCol); - const dy = Math.abs(imgY - activeRoiCenterRow); + const geometry = resizeDetectorFromPointer({ + mode: roiMode as DetectorRoiMode, + centerRow: activeRoiCenterRow, + centerCol: activeRoiCenterCol, + pointerRow: imgY, + pointerCol: imgX, + radius: roiRadius, + radiusInner: roiRadiusInner, + aspectRatio: resizeAspectRef.current, + preserveAspect: shiftKey, + }); + if (!geometry) return false; if (roiMode === "rect") { - let newW = Math.max(2, Math.round(dx * 2)); - let newH = Math.max(2, Math.round(dy * 2)); - if (shiftKey && resizeAspectRef.current != null) { - const aspect = resizeAspectRef.current; - if (newW / newH > aspect) newH = Math.max(2, Math.round(newW / aspect)); - else newW = Math.max(2, Math.round(newH * aspect)); - } - setRoiWidth(newW); - setRoiHeight(newH); + setRoiWidth(geometry.width!); + setRoiHeight(geometry.height!); } else { - const newRadius = roiMode === "square" ? Math.max(dx, dy) : Math.sqrt(dx ** 2 + dy ** 2); - const minRadius = roiMode === "annular" ? (roiRadiusInner || 0) + 1 : 1; - const rad = Math.max(minRadius, Math.round(newRadius)); + const rad = geometry.radius!; setLocalRoiRadius(rad); sendRoiRadius(rad); } @@ -8543,8 +8628,12 @@ function Show4DSTEM() { dpDragOffsetRef.current = { dRow: 0, dCol: 0 }; setLocalKCol(imgX); setLocalKRow(imgY); // Use compound roi_center trait [row, col] - single observer fires in Python - const newCol = Math.round(Math.max(0, Math.min(detCols - 1, imgX))); - const newRow = Math.round(Math.max(0, Math.min(detRows - 1, imgY))); + const { row: newRow, col: newCol } = clampDetectorCenter( + imgY, + imgX, + detRows, + detCols, + ); model.set("roi_active", true); writeRoiCenterModel(newRow, newCol); requestCompareViLive(); @@ -8648,8 +8737,16 @@ function Show4DSTEM() { const centerRow = imgY - dpDragOffsetRef.current.dRow; setLocalKCol(centerCol); setLocalKRow(centerRow); // rAF-coalesced — sends only the latest roi_center per frame. - const newCol = Math.round(Math.max(0, Math.min(detCols - 1, centerCol))); - const newRow = Math.round(Math.max(0, Math.min(detRows - 1, centerRow))); + // Keep the detector geometry subpixel while dragging. The public traits are + // floats and the scientific mask evaluates detector-pixel centers against + // that geometry. Rounding here made a binned 48x48 detector update only + // every roughly ten screen pixels, which looked like a pointer-up commit. + const { row: newRow, col: newCol } = clampDetectorCenter( + centerRow, + centerCol, + detRows, + detCols, + ); queueRoiCenter(newRow, newCol); requestCompareViLive(); }; @@ -10568,6 +10665,7 @@ function Show4DSTEM() { count={comparePanelCount || 0} indices={comparePanelIndices || []} gpuSlots={compareGpuSlotsRef.current} + gpuRanges={compareGpuRangesRef.current} gpuVersion={compareGpuVersion} gpuEngine={viGpuColormapRef.current} progressivePage={progressiveComparePage} diff --git a/scripts/sync-gpu-webgpu.mjs b/scripts/sync-gpu-webgpu.mjs index 5d83c488..d132ed37 100644 --- a/scripts/sync-gpu-webgpu.mjs +++ b/scripts/sync-gpu-webgpu.mjs @@ -31,11 +31,15 @@ names = ( "display/webgpu/stats.ts", "display/goldens/parity.json", "swift/Sources/MetalDisplayKernels/Resources/colormaps.json", + "parity/scan_rotation_v1.json", + "geometry/compute/webgpu/quarter-turn.ts", "io/backends/webgpu/bslz4.ts", "io/backends/webgpu/h5reader.ts", + "io/backends/webgpu/logical-pixel-hash.ts", "io/backends/webgpu/local-h5.ts", - "detector/compute/webgpu/backend.ts", "detector/geometry.ts", + "detector/compute/webgpu/exact-com.ts", + "detector/compute/webgpu/backend.ts", "dpc/compute/webgpu/fft.ts", "dpc/compute/webgpu/kernels.ts", "ssb/compute/webgpu/backend.ts", diff --git a/src/quantem/widget/cli.py b/src/quantem/widget/cli.py index a14dc62d..b6711770 100644 --- a/src/quantem/widget/cli.py +++ b/src/quantem/widget/cli.py @@ -19,6 +19,7 @@ instead of preprocessing every frame before the viewer opens. """ import argparse +import copy import email.utils import http.server import json @@ -31,6 +32,7 @@ import socketserver import sys import threading +import tempfile import urllib.parse import webbrowser @@ -472,6 +474,8 @@ def _render_html(args: argparse.Namespace) -> int: # --------------------------------------------------------------------------- _WIDGET_CELL = ("Show2D(", "Show3D(", "Show4DSTEM(", "Show3DSlices(", "ShowEDS(") +_WIDGET_STATE_MIME = "application/vnd.jupyter.widget-state+json" +_WIDGET_VIEW_MIME = "application/vnd.jupyter.widget-view+json" def _add_github_args(parser: argparse.ArgumentParser) -> None: @@ -481,6 +485,8 @@ def _add_github_args(parser: argparse.ArgumentParser) -> None: help="Use the notebook's existing outputs instead of re-running it.") parser.add_argument("--quality", type=int, default=92, help="JPEG quality for the embedded renders (default 92).") + parser.add_argument("--max-width", type=int, default=1200, + help="Maximum embedded UI width in pixels (default 1200).") parser.add_argument("--timeout", type=int, default=600, help="Per-cell execution timeout in seconds (default 600).") @@ -489,11 +495,24 @@ def _strip_state(nb: dict) -> None: """Drop the heavy offline live-widget manager-state + the dead widget-view output refs.""" nb.get("metadata", {}).pop("widgets", None) for cell in nb.get("cells", []): + kept = [] for out in cell.get("outputs", []): (out.get("data") or {}).pop("application/vnd.jupyter.widget-view+json", None) - - -def _embed_jpeg(cell: dict, png_or_jpeg: bytes, quality: int) -> bool: + if out.get("output_type") in {"display_data", "execute_result"} and not ( + out.get("data") or {} + ): + continue + kept.append(out) + if "outputs" in cell: + cell["outputs"] = kept + + +def _embed_jpeg( + cell: dict, + png_or_jpeg: bytes, + quality: int, + max_width: int = 1200, +) -> bool: """Replace a cell's visual output with one JPEG. Widget outputs usually have only ``application/vnd.jupyter.widget-view+json``, @@ -504,6 +523,9 @@ def _embed_jpeg(cell: dict, png_or_jpeg: bytes, quality: int) -> bool: from io import BytesIO from PIL import Image img = Image.open(BytesIO(png_or_jpeg)).convert("RGB") + if max_width > 0 and img.width > max_width: + height = max(1, round(img.height * max_width / img.width)) + img = img.resize((max_width, height), Image.Resampling.LANCZOS) buf = BytesIO() img.save(buf, format="JPEG", quality=quality, optimize=True) b64 = base64.b64encode(buf.getvalue()).decode("ascii") @@ -517,13 +539,21 @@ def _embed_jpeg(cell: dict, png_or_jpeg: bytes, quality: int) -> bool: for k in [k for k in data if k.startswith("image/")]: del data[k] data["image/jpeg"] = b64 - out.setdefault("metadata", {}) + metadata = out.setdefault("metadata", {}) + quantem_metadata = metadata.setdefault("quantem.widget", {}) + quantem_metadata["github_full_ui"] = True + quantem_metadata["github_quality"] = quality + quantem_metadata["github_width"] = img.width done = True break if not done: cell.setdefault("outputs", []).append({ "output_type": "display_data", - "metadata": {}, + "metadata": {"quantem.widget": { + "github_full_ui": True, + "github_quality": quality, + "github_width": img.width, + }}, "data": {"image/jpeg": b64}, }) done = True @@ -548,6 +578,312 @@ def _cell_has_widget_view_output(cell: dict) -> bool: return False +def _cell_has_full_ui_output(cell: dict) -> bool: + """Return true when ``quantem github`` already embedded the full widget UI.""" + for out in cell.get("outputs", []): + data = out.get("data") or {} + metadata = out.get("metadata") or {} + quantem_metadata = metadata.get("quantem.widget") or {} + if any(key.startswith("image/") for key in data) and quantem_metadata.get( + "github_full_ui" + ) is True: + return True + return False + + +def _github_widget_cells(nb: dict) -> list[dict]: + """Find widget cells from runtime output first, with source as a fallback. + + Public APIs such as ``drift.show()`` return QuantEM widgets without naming + ``Show2D`` in notebook source. Runtime widget MIME is therefore the + authoritative signal after execution. The source check retains support + for older or hand-edited notebooks, and the metadata marker recognizes a + notebook already prepared by this command. + """ + return [ + cell + for cell in nb.get("cells", []) + if cell.get("cell_type") == "code" + and ( + _cell_has_widget_view_output(cell) + or _cell_has_full_ui_output(cell) + or any(widget in "".join(cell.get("source", [])) for widget in _WIDGET_CELL) + ) + ] + + +def _github_capture_cells(nb: dict) -> list[dict]: + """Return widget cells that still need a browser-captured full-UI image.""" + return [ + cell + for cell in _github_widget_cells(nb) + if not _cell_has_full_ui_output(cell) + and ( + _cell_has_widget_view_output(cell) + or not _cell_has_image_output(cell) + ) + ] + + +def _widget_view_model_ids(cell: dict) -> list[str]: + """Return root widget model IDs referenced by one output cell.""" + model_ids = [] + for out in cell.get("outputs", []): + view = (out.get("data") or {}).get(_WIDGET_VIEW_MIME) + model_id = view.get("model_id") if isinstance(view, dict) else None + if model_id and model_id not in model_ids: + model_ids.append(model_id) + return model_ids + + +def _widget_model_closure(state: dict, roots: list[str]) -> set[str]: + """Return each root model and every ``IPY_MODEL_`` dependency it references.""" + found: set[str] = set() + pending = list(roots) + + def references(value): + if isinstance(value, str) and value.startswith("IPY_MODEL_"): + yield value.removeprefix("IPY_MODEL_") + elif isinstance(value, dict): + for item in value.values(): + yield from references(item) + elif isinstance(value, list): + for item in value: + yield from references(item) + + while pending: + model_id = pending.pop() + if model_id in found: + continue + if model_id not in state: + raise ValueError(f"widget model {model_id!r} is absent from notebook state") + found.add(model_id) + pending.extend(ref for ref in references(state[model_id]) if ref not in found) + return found + + +def _widget_capture_notebook(nb: dict, cell: dict) -> dict: + """Build a minimal notebook containing one live widget and its dependencies. + + A scientific notebook can contain hundreds of megabytes of state per widget. + Rendering every model into one HTML document can exceed the browser's JSON + parser limit even though each widget renders correctly on its own. This + temporary notebook keeps exactly the state required by one output view. + """ + widget_payload = ( + nb.get("metadata", {}).get("widgets", {}).get(_WIDGET_STATE_MIME) + ) + state = widget_payload.get("state") if isinstance(widget_payload, dict) else None + roots = _widget_view_model_ids(cell) + if not roots: + raise ValueError("widget output has no model_id to capture") + if not isinstance(state, dict): + raise ValueError("notebook has no saved widget state; execute it before capture") + keep = _widget_model_closure(state, roots) + + capture_cell = copy.deepcopy(cell) + capture_cell["source"] = [] + capture_cell["outputs"] = [ + { + "output_type": out.get("output_type", "display_data"), + "metadata": {}, + "data": {_WIDGET_VIEW_MIME: copy.deepcopy((out.get("data") or {})[_WIDGET_VIEW_MIME])}, + } + for out in cell.get("outputs", []) + if _WIDGET_VIEW_MIME in (out.get("data") or {}) + ] + metadata = { + key: copy.deepcopy(value) + for key, value in nb.get("metadata", {}).items() + if key != "widgets" + } + capture_payload = { + key: copy.deepcopy(value) + for key, value in widget_payload.items() + if key != "state" + } + capture_payload["state"] = { + model_id: copy.deepcopy(state[model_id]) for model_id in keep + } + metadata["widgets"] = {_WIDGET_STATE_MIME: capture_payload} + return { + "cells": [capture_cell], + "metadata": metadata, + "nbformat": nb.get("nbformat", 4), + "nbformat_minor": nb.get("nbformat_minor", 5), + } + + +def _capture_notebook_widget_uis( + notebook: pathlib.Path, + nb: dict, + capture_cells: list[dict], +) -> list[bytes]: + """Render and capture widget cells independently to bound temporary HTML size.""" + import subprocess + + shots: list[bytes] = [] + with tempfile.TemporaryDirectory( + prefix=f".{notebook.stem}-github-ui-", dir=notebook.parent + ) as folder: + temporary = pathlib.Path(folder) + for index, cell in enumerate(capture_cells, start=1): + capture_nb = temporary / f"widget-{index:02d}.ipynb" + capture_nb.write_text( + json.dumps(_widget_capture_notebook(nb, cell)), encoding="utf-8" + ) + result = subprocess.run( + [ + "jupyter", "nbconvert", "--to", "html", str(capture_nb), + "--output-dir", str(temporary), "--output", capture_nb.stem, + ] + ) + if result.returncode != 0: + raise ValueError( + f"nbconvert failed while preparing widget UI {index}" + ) + html = temporary / f"{capture_nb.stem}.html" + print( + f" widget {index}/{len(capture_cells)} temporary HTML: " + f"{html.stat().st_size / 1e6:.1f} MB" + ) + captured = _capture_full_ui(html, 1) + if len(captured) != 1: + raise ValueError( + f"captured {len(captured)} UI screenshot(s) for widget {index}" + ) + shots.extend(captured) + return shots + + +def _recompress_full_ui_outputs(nb: dict, quality: int, max_width: int) -> int: + """Re-encode previously prepared full-UI images at the requested quality.""" + import base64 + + changed = 0 + for cell in nb.get("cells", []): + for out in cell.get("outputs", []): + metadata = (out.get("metadata") or {}).get("quantem.widget") or {} + data = out.get("data") or {} + image = data.get("image/jpeg") + if ( + metadata.get("github_full_ui") is True + and image + and ( + metadata.get("github_quality") != quality + or metadata.get("github_width") != min( + max_width, metadata.get("github_width", max_width + 1) + ) + ) + ): + _embed_jpeg(cell, base64.b64decode(image), quality, max_width) + changed += 1 + break + return changed + + +def _prune_widget_fallbacks(nb: dict) -> int: + """Remove redundant auto-snapshots after a complete UI capture exists.""" + removed = 0 + for cell in _github_widget_cells(nb): + if not _cell_has_full_ui_output(cell): + continue + kept = [] + for out in cell.get("outputs", []): + quantem_metadata = (out.get("metadata") or {}).get("quantem.widget") or {} + if ( + quantem_metadata.get("static_fallback") is True + and quantem_metadata.get("github_full_ui") is not True + ): + removed += 1 + continue + if quantem_metadata.get("github_full_ui") is True: + data = out.get("data") or {} + image = data.get("image/jpeg") + out["data"] = {"image/jpeg": image} if image else {} + kept.append(out) + cell["outputs"] = kept + return removed + + +def _validate_github_widget_outputs(widget_cells: list[dict]) -> None: + """Require one browser-captured UI and no duplicate widget render per cell.""" + + problems = [] + for index, cell in enumerate(widget_cells, start=1): + full_ui = [] + fallbacks = 0 + widget_views = 0 + for out in cell.get("outputs", []): + data = out.get("data") or {} + quantem_metadata = (out.get("metadata") or {}).get( + "quantem.widget" + ) or {} + if quantem_metadata.get("github_full_ui") is True and any( + key.startswith("image/") for key in data + ): + full_ui.append(out) + if quantem_metadata.get("static_fallback") is True: + fallbacks += 1 + if _WIDGET_VIEW_MIME in data: + widget_views += 1 + if len(full_ui) != 1 or fallbacks or widget_views: + problems.append( + f"cell {index}: full_ui={len(full_ui)}, " + f"fallbacks={fallbacks}, widget_views={widget_views}" + ) + if problems: + raise ValueError( + "GitHub notebook preparation requires exactly one browser-captured " + "widget UI per widget cell and no fallback duplicates: " + + "; ".join(problems) + ) + + +def _compress_large_raster_outputs( + nb: dict, + quality: int, + max_width: int, + threshold: int = 500_000, +) -> int: + """JPEG-encode large ordinary PNG outputs while retaining readable dimensions.""" + import base64 + from io import BytesIO + from PIL import Image + + changed = 0 + for cell in nb.get("cells", []): + for out in cell.get("outputs", []): + metadata = (out.get("metadata") or {}).get("quantem.widget") or {} + if metadata.get("github_full_ui") is True: + continue + data = out.get("data") or {} + encoded = data.get("image/png") + if not encoded or len(encoded) <= threshold: + continue + image = Image.open(BytesIO(base64.b64decode(encoded))) + if image.mode in {"RGBA", "LA"}: + rgba = image.convert("RGBA") + background = Image.new("RGBA", rgba.size, "white") + background.alpha_composite(rgba) + image = background.convert("RGB") + else: + image = image.convert("RGB") + if max_width > 0 and image.width > max_width: + height = max(1, round(image.height * max_width / image.width)) + image = image.resize((max_width, height), Image.Resampling.LANCZOS) + buffer = BytesIO() + image.save(buffer, format="JPEG", quality=quality, optimize=True) + data.pop("image/png") + data["image/jpeg"] = base64.b64encode(buffer.getvalue()).decode("ascii") + metadata = out.setdefault("metadata", {}).setdefault("quantem.widget", {}) + metadata["github_compressed_from"] = "image/png" + metadata["github_quality"] = quality + metadata["github_width"] = image.width + changed += 1 + return changed + + def _capture_full_ui(html: pathlib.Path, n_expected: int) -> list[bytes]: """Screenshot each widget's FULL UI (toolbar + toggles + panels + histograms) from the rendered live-widget HTML, deterministically, via Playwright on the real GPU. The widget @@ -569,8 +905,30 @@ def _capture_full_ui(html: pathlib.Path, n_expected: int) -> list[bytes]: "--enable-unsafe-webgpu", "--use-angle=vulkan", "--enable-features=Vulkan", "--ignore-gpu-blocklist", "--disable-gpu-sandbox", "--no-sandbox"], **launch_kwargs) page = browser.new_page(viewport={"width": 1300, "height": 2400}, device_scale_factor=2) - page.goto(html.as_uri(), wait_until="load", timeout=90000) - page.wait_for_timeout(13000) # anywidget mount + WebGPU paint + browser_errors = [] + page.on("pageerror", lambda error: browser_errors.append(f"page: {error}")) + page.on( + "console", + lambda message: browser_errors.append(f"console: {message.text}") + if message.type == "error" + else None, + ) + page.goto(html.as_uri(), wait_until="load", timeout=180000) + # Large scientific notebooks can carry hundreds of megabytes of + # temporary widget state. Wait for actual canvases instead of assuming + # that every browser mounts and paints them within a fixed 13 seconds. + canvas_count = 0 + for _ in range(120): + canvas_count = page.locator(".jp-OutputArea-output canvas").count() + if canvas_count >= n_expected: + break + page.wait_for_timeout(1000) + if canvas_count < n_expected: + for error in browser_errors[-10:]: + print(f" browser error: {error}") + print(f" mounted canvases: {canvas_count}/{n_expected}") + else: + page.wait_for_timeout(2000) # allow the first WebGPU frame to present arch = page.evaluate("async()=>{const a=await navigator.gpu?.requestAdapter();" "return a?(a.info?.architecture||'?'):'none';}") print(f" GPU adapter: {arch}") @@ -613,30 +971,20 @@ def _prepare_github(args: argparse.Namespace) -> int: if not args.no_execute: print(f"executing {notebook.name} ...") if subprocess.run(["jupyter", "nbconvert", "--to", "notebook", "--execute", "--inplace", - str(notebook), f"--ExecutePreprocessor.timeout={args.timeout}"]).returncode != 0: + str(notebook), f"--ExecutePreprocessor.timeout={args.timeout}", + "--ExecutePreprocessor.store_widget_state=True"]).returncode != 0: raise ValueError("nbconvert --execute failed (see output above)") nb = json.loads(notebook.read_text()) - widget_cells = [c for c in nb["cells"] - if c["cell_type"] == "code" and any(w in "".join(c["source"]) for w in _WIDGET_CELL)] - capture_cells = [ - c for c in widget_cells - if not _cell_has_image_output(c) - ] + widget_cells = _github_widget_cells(nb) + capture_cells = _github_capture_cells(nb) + max_width = getattr(args, "max_width", 1200) + recompressed = _recompress_full_ui_outputs(nb, args.quality, max_width) if capture_cells: try: - html = notebook.with_suffix(".fullui.html") - subprocess.run(["jupyter", "nbconvert", "--to", "html", str(notebook), - "--output-dir", str(notebook.parent), "--output", notebook.stem + ".fullui"], - check=True) print(f"capturing {len(capture_cells)} widget UI(s) on the GPU ...") - shots = _capture_full_ui(html, len(capture_cells)) - html.unlink(missing_ok=True) - if len(shots) != len(capture_cells): - raise ValueError( - f"captured {len(shots)} widget UI screenshot(s) for {len(capture_cells)} widget cell(s)" - ) + shots = _capture_notebook_widget_uis(notebook, nb, capture_cells) for cell, png in zip(capture_cells, shots): - _embed_jpeg(cell, png, args.quality) + _embed_jpeg(cell, png, args.quality, max_width) mode = f"{len(shots)} full-UI screenshots" except (ImportError, RuntimeError, OSError) as err: raise ValueError( @@ -646,11 +994,20 @@ def _prepare_github(args: argparse.Namespace) -> int: mode = f"{len(widget_cells)} existing image output(s)" else: mode = "no widget cells - state stripped only" + fallbacks = _prune_widget_fallbacks(nb) + rasters = _compress_large_raster_outputs(nb, args.quality, max_width) _strip_state(nb) + _validate_github_widget_outputs(widget_cells) notebook.write_text(json.dumps(nb, indent=1)) after = notebook.stat().st_size print(f"github-ready: {notebook.name} {before / 1e6:.1f} MB -> {after / 1e6:.1f} MB" - f" ({mode}, JPEG q{args.quality}, offline state stripped)") + f" ({mode}, JPEG q{args.quality}, max {max_width}px, offline state stripped)") + if recompressed: + print(f" re-encoded {recompressed} existing full-UI image(s) at JPEG q{args.quality}") + if fallbacks: + print(f" removed {fallbacks} redundant widget fallback image(s)") + if rasters: + print(f" compressed {rasters} large raster output(s) for repository display") if after > 5e6: print(" warning: still > 5 MB - GitHub may not render. Lower --quality or the widget's size=.") return 0 diff --git a/src/quantem/widget/image_folder.py b/src/quantem/widget/image_folder.py index 928779bc..0c5826a7 100644 --- a/src/quantem/widget/image_folder.py +++ b/src/quantem/widget/image_folder.py @@ -448,6 +448,12 @@ def _poll_once(self, widget: Any) -> list[int]: f"Applying {len(changed)} stable image file" f"{'s' if len(changed) != 1 else ''}.", ) + # Publish the first-frame transition before replacing the resident + # array. ``set_image`` exposes its new frame count partway through the + # update, so a concurrent notebook/UI reader must never observe + # ``n_slices > 0`` while ``folder_waiting`` is still true. + _safe_set_widget_status(widget, "_folder_waiting", False) + _safe_set_widget_status(widget, "folder_waiting", False) widget._apply_folder_image_records( old_records, new_records, diff --git a/src/quantem/widget/info.py b/src/quantem/widget/info.py index 34e28f54..bfe8ce77 100644 --- a/src/quantem/widget/info.py +++ b/src/quantem/widget/info.py @@ -3,6 +3,7 @@ import json import os import platform +import re import subprocess from datetime import datetime from importlib.metadata import PackageNotFoundError, distribution, version @@ -14,6 +15,13 @@ from packaging.version import Version +def _concise_cuda_name(name: str) -> str: + """Return a stable, readable CUDA device label for notebook reports.""" + + match = re.match(r"^(NVIDIA RTX PRO \d+)", str(name).strip()) + return match.group(1) if match else str(name).strip() + + def profile(*, check_updates: bool = False) -> None: """Print the installed QuantEM stack and active compute environment. @@ -126,7 +134,7 @@ def print_distribution_status( import torch if torch.cuda.is_available(): - device = f"cuda ({torch.cuda.get_device_name(0)})" + device = f"cuda ({_concise_cuda_name(torch.cuda.get_device_name(0))})" elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): device = "mps (Apple)" else: diff --git a/tests/show4dstem/test_vi_kernel_contract.py b/tests/show4dstem/test_vi_kernel_contract.py index 42c31151..c046d066 100644 --- a/tests/show4dstem/test_vi_kernel_contract.py +++ b/tests/show4dstem/test_vi_kernel_contract.py @@ -114,6 +114,9 @@ def test_show4dstem_webgpu_engine_has_selected_index_vi_kernel() -> None: assert "renderSlotDirectWithGpuRangeToCanvas" in ( repo / "js" / ".generated" / "engine" / "display" / "webgpu" / "colormaps.ts" ).read_text(encoding="utf-8") + assert "renderPanelSlotsToImageBitmapAsync" in ( + repo / "js" / ".generated" / "engine" / "display" / "webgpu" / "colormaps.ts" + ).read_text(encoding="utf-8") assert "function buildDetectorMask" not in frontend assert "function buildScanMask" not in frontend assert "buildFullDetectorMask" in frontend @@ -135,6 +138,11 @@ def test_show4dstem_webgpu_engine_has_selected_index_vi_kernel() -> None: assert "virtualGpuCanvasRef" not in frontend assert "renderPanelSlotsDirectToCanvas" not in frontend assert "renderSlotDirectWithGpuRangeToCanvas" in frontend + assert "compareGpuRangesRef" in frontend + assert "computeRangeBatch(batchSlots)" in frontend + assert "gpuRanges={compareGpuRangesRef.current}" in frontend + assert "rangeReadbackBytes" in frontend + assert "gpuOnlyHotPath: stats.lastRangeReadbackBytes === 0" in frontend def test_show4dstem_webgpu_h5_master_loader_batches_external_decodes() -> None: @@ -376,22 +384,61 @@ def test_show4dstem_multiple_detector_drag_uses_live_gpu_compare_slots() -> None "requestViFinalizeRef.current", 1, )[0] - assert "await recomputeVI();" in live_drag - assert "await recomputeCompareVI();" in live_drag - assert live_drag.index("await recomputeVI();") < live_drag.index( - "await recomputeCompareVI();" + visible_route = frontend.split( + "const recomputeVisibleVirtualImages = async () => {", + 1, + )[1].split( + '(window as unknown as { __sh4d: unknown })', + 1, + )[0] + assert 'mode === "multiple" || mode === "compare"' in visible_route + assert visible_route.index("await recomputeCompareVI();") < visible_route.index( + "await recomputeVI();" ) + assert "void recomputeVisibleVirtualImages().finally" in live_drag + assert "recomputeVI" not in live_drag + assert "recomputeCompareVI" not in live_drag + finalize = frontend.split("requestViFinalizeRef.current = () => {", 1)[1].split( + "};", + 1, + )[0] + assert "void recomputeVisibleVirtualImages();" in finalize + detector_drag = frontend.split( + "const handleDpMouseMove =", + 1, + )[1].split( + "const handleDpMouseUp =", + 1, + )[0] + assert "Keep the detector geometry subpixel while dragging." in detector_drag + assert "Math.round(Math.max(0, Math.min(detCols - 1, centerCol)))" not in ( + detector_drag + ) + detector_resize = frontend.split( + "const resizeDpRoiFromImagePoint =", + 1, + )[1].split( + "React.useEffect(() => {", + 1, + )[0] + assert "Math.round(newRadius)" not in detector_resize assert 'type DpcGpuSource = "DPC_row" | "DPC_col" | "iDPC";' in frontend - assert "gpuLoaded: Boolean(gpuSlots?.has(frame) && gpuEngine)" in frontend + assert "gpuLoaded: Boolean(gpuSlots?.has(frame) && gpuRanges?.has(frame) && gpuEngine)" in frontend assert 'scaleMode === "log"' in frontend assert "entry.panel !== undefined || entry.gpuLoaded" in frontend assert "const loaded = panel !== undefined || gpuLoaded;" in frontend assert "onChangeCommitted={finishDpRoiInteraction}" in frontend assert "__sh4dLiveViStats" in frontend - assert "gpuOnlyHotPath: true" in frontend + assert "gpuOnlyHotPath: stats.lastRangeReadbackBytes === 0" in frontend assert 'publishLiveCompareViStats("paint"' in frontend assert "if (gpuEngine) gpuEngine.uploadLUT(colormap, lut);" in frontend - assert "renderSlotDirectWithGpuRangeToCanvas" in frontend + assert "renderPanelSlotsToImageBitmapAsync" in frontend + assert "width: shapeCols * panels.length" in frontend + assert "panelCount: panels.length" in frontend + assert "cols: panels.length" in frontend + assert "index * shapeCols" in frontend + assert 'panel.canvas.getContext("2d")' in frontend + assert "computeRangeBatch(batchSlots)" in frontend assert "renderSlotGpuRangeToOffscreen" not in frontend assert "let comparePersistentStack: Float32Array | null = null;" in frontend assert "if (getVol && !volIsResident(idx)) continue;" in frontend diff --git a/tests/show4dstem/test_webgpu_bundle.py b/tests/show4dstem/test_webgpu_bundle.py index 9efe4737..90052a17 100644 --- a/tests/show4dstem/test_webgpu_bundle.py +++ b/tests/show4dstem/test_webgpu_bundle.py @@ -74,6 +74,7 @@ def test_bundle_export_writes_launcher_viewer_and_vendored_page(tmp_path): assert "__QT_REQUIRE_LOCAL_H5_FILES" in page assert "__QT_H5_DECODE_DTYPE" in page and "__BSLZ4_FRAME_WG" in page assert 'globalThis.__QT_H5_DECODE_DTYPE ??= "u2"' in page + assert "globalThis.__QT_H5_MAX_RESIDENT ??=" not in page assert "globalThis.__QT_H5_FORCE_LOW8 ??= false" in page assert "globalThis.__BSLZ4_PIPELINE_STAGING ??= false" in page assert "../tilt_a_master.h5" in page @@ -98,6 +99,7 @@ def test_bundle_export_uses_low8_for_audited_uint8_h5(tmp_path): widget.close() page = (tmp_path / ".viewer" / "Show4DSTEM.html").read_text(encoding="utf-8") assert 'globalThis.__QT_H5_DECODE_DTYPE ??= "uint8"' in page + assert "globalThis.__QT_H5_MAX_RESIDENT ??=" not in page assert "globalThis.__QT_H5_FORCE_LOW8 ??= true" in page assert "globalThis.__BSLZ4_LOW8_ONLY ??= true" in page assert "globalThis.__QT_H5_MAX_RESIDENT ??= 1" not in page diff --git a/tests/test_cli.py b/tests/test_cli.py index de8b846e..8a4f2058 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -38,16 +38,150 @@ def test_embed_jpeg_adds_image_to_widget_only_output(tmp_path): } assert cli._embed_jpeg(cell, png.read_bytes(), quality=80) - data = cell["outputs"][0]["data"] + output = cell["outputs"][0] + data = output["data"] assert "image/jpeg" in data assert "application/vnd.jupyter.widget-view+json" in data + assert output["metadata"]["quantem.widget"]["github_full_ui"] is True + assert output["metadata"]["quantem.widget"]["github_quality"] == 80 + assert output["metadata"]["quantem.widget"]["github_width"] == 24 def test_github_widget_cell_detector_includes_showeds(): assert "ShowEDS(" in cli._WIDGET_CELL -def test_github_prepare_reuses_existing_image_outputs(tmp_path, monkeypatch): +def test_github_widget_cell_detector_uses_runtime_widget_output_for_public_api(): + cell = { + "cell_type": "code", + "source": ["drift.show(mode='interactive')"], + "outputs": [{ + "output_type": "display_data", + "metadata": {}, + "data": { + "application/vnd.jupyter.widget-view+json": {"model_id": "abc"}, + "image/jpeg": "fallback", + }, + }], + } + notebook = {"cells": [cell]} + + assert cli._github_widget_cells(notebook) == [cell] + assert cli._github_capture_cells(notebook) == [cell] + + +def test_github_capture_reuses_only_marked_full_ui_output(): + cell = { + "cell_type": "code", + "source": ["drift.show(mode='interactive')"], + "outputs": [{ + "output_type": "display_data", + "metadata": {"quantem.widget": {"github_full_ui": True}}, + "data": {"image/jpeg": "full-ui"}, + }], + } + notebook = {"cells": [cell]} + + assert cli._github_widget_cells(notebook) == [cell] + assert cli._github_capture_cells(notebook) == [] + + +def test_widget_model_closure_includes_layout_dependency(): + state = { + "root": {"state": {"layout": "IPY_MODEL_layout"}}, + "layout": {"state": {}}, + "unrelated": {"state": {}}, + } + + assert cli._widget_model_closure(state, ["root"]) == {"root", "layout"} + + +def test_widget_capture_notebook_keeps_only_required_models(): + view = {"model_id": "root", "version_major": 2, "version_minor": 0} + cell = { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": ["drift.show()"], + "outputs": [{ + "output_type": "display_data", + "metadata": {}, + "data": {cli._WIDGET_VIEW_MIME: view, "image/jpeg": "fallback"}, + }], + } + notebook = { + "cells": [cell], + "metadata": {"widgets": {cli._WIDGET_STATE_MIME: { + "version_major": 2, + "version_minor": 0, + "state": { + "root": {"state": {"layout": "IPY_MODEL_layout"}}, + "layout": {"state": {}}, + "unrelated": {"state": {}}, + }, + }}}, + "nbformat": 4, + "nbformat_minor": 5, + } + + capture = cli._widget_capture_notebook(notebook, cell) + payload = capture["metadata"]["widgets"][cli._WIDGET_STATE_MIME] + assert set(payload["state"]) == {"root", "layout"} + assert capture["cells"][0]["source"] == [] + assert capture["cells"][0]["outputs"][0]["data"] == { + cli._WIDGET_VIEW_MIME: view + } + + +def test_prune_widget_fallbacks_keeps_only_full_ui_visual(): + cell = { + "cell_type": "code", + "source": ["drift.show()"], + "outputs": [ + { + "output_type": "display_data", + "metadata": {"quantem.widget": {"github_full_ui": True}}, + "data": {"image/jpeg": "ui", "text/html": "redundant"}, + }, + { + "output_type": "display_data", + "metadata": {"quantem.widget": {"static_fallback": True}}, + "data": {"image/jpeg": "fallback", "text/html": "fallback"}, + }, + ], + } + + assert cli._prune_widget_fallbacks({"cells": [cell]}) == 1 + assert len(cell["outputs"]) == 1 + assert cell["outputs"][0]["data"] == {"image/jpeg": "ui"} + + +def test_github_validation_rejects_duplicate_fallback(): + cell = { + "cell_type": "code", + "source": ["drift.show()"], + "outputs": [ + { + "output_type": "display_data", + "metadata": {"quantem.widget": {"github_full_ui": True}}, + "data": {"image/jpeg": "ui"}, + }, + { + "output_type": "display_data", + "metadata": {"quantem.widget": {"static_fallback": True}}, + "data": {"image/jpeg": "fallback"}, + }, + ], + } + + with pytest.raises(ValueError, match="fallbacks=1"): + cli._validate_github_widget_outputs([cell]) + + assert cli._prune_widget_fallbacks({"cells": [cell]}) == 1 + cli._validate_github_widget_outputs([cell]) + + +def test_github_prepare_reuses_existing_full_ui_output(tmp_path, monkeypatch): notebook = tmp_path / "show2d_github.ipynb" notebook.write_text( """{ @@ -59,7 +193,13 @@ def test_github_prepare_reuses_existing_image_outputs(tmp_path, monkeypatch): "outputs": [ { "output_type": "display_data", - "metadata": {}, + "metadata": { + "quantem.widget": { + "github_full_ui": true, + "github_quality": 90, + "github_width": 1200 + } + }, "data": { "text/plain": "", "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/2w==" @@ -85,7 +225,7 @@ def test_github_prepare_reuses_existing_image_outputs(tmp_path, monkeypatch): ) def fail_capture(*args, **kwargs): - raise AssertionError("existing image outputs should not trigger browser capture") + raise AssertionError("existing full-UI output should not trigger capture") monkeypatch.setattr(cli, "_capture_full_ui", fail_capture) args = type("Args", (), { @@ -98,6 +238,7 @@ def fail_capture(*args, **kwargs): assert cli._prepare_github(args) == 0 text = notebook.read_text(encoding="utf-8") assert "image/jpeg" in text + assert text.count("github_full_ui") == 1 assert "application/vnd.jupyter.widget-state+json" not in text diff --git a/tests/test_profile.py b/tests/test_profile.py index 6eaa19e1..03bce68d 100644 --- a/tests/test_profile.py +++ b/tests/test_profile.py @@ -4,6 +4,19 @@ from types import SimpleNamespace +def test_profile_shortens_rtx_pro_marketing_name() -> None: + """Notebook profiles show the GPU model without workstation marketing text.""" + from quantem.widget.info import _concise_cuda_name + + assert ( + _concise_cuda_name( + "NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition" + ) + == "NVIDIA RTX PRO 6000" + ) + assert _concise_cuda_name("NVIDIA H100 80GB HBM3") == "NVIDIA H100 80GB HBM3" + + def test_profile_reports_the_installed_quantem_stack(capsys) -> None: """A notebook records every QuantEM package through one profile call.""" import quantem.widget as qw diff --git a/tests/test_quantem_gpu_ownership.py b/tests/test_quantem_gpu_ownership.py index 5344f1f4..a32484c1 100644 --- a/tests/test_quantem_gpu_ownership.py +++ b/tests/test_quantem_gpu_ownership.py @@ -354,7 +354,12 @@ def test_widget_webgpu_sources_are_generated_from_quantem_gpu() -> None: assert '"display/webgpu/geometry.ts"' in sync_script assert '"display/webgpu/stats.ts"' in sync_script assert '"swift/Sources/MetalDisplayKernels/Resources/colormaps.json"' in sync_script + assert '"parity/scan_rotation_v1.json"' in sync_script + assert '"geometry/compute/webgpu/quarter-turn.ts"' in sync_script assert '"io/backends/webgpu/bslz4.ts"' in sync_script + assert '"io/backends/webgpu/logical-pixel-hash.ts"' in sync_script + assert '"detector/geometry.ts"' in sync_script + assert '"detector/compute/webgpu/exact-com.ts"' in sync_script assert '"detector/compute/webgpu/backend.ts"' in sync_script assert '"dpc/compute/webgpu/fft.ts"' in sync_script assert "syncGpuWebgpuSources()" in build_script