Skip to content

Latest commit

 

History

History
1428 lines (1025 loc) · 48.3 KB

File metadata and controls

1428 lines (1025 loc) · 48.3 KB

Developing quantem.widget

Welcome! This guide walks you through setting up, understanding, and contributing to quantem.widget - interactive Jupyter widgets for electron microscopy. Each widget has a Python backend (data processing, GPU computation) and a JavaScript frontend (rendering, user interaction), connected by anywidget.

For contributing to the core quantem library, see ophusgroup/dev.

Table of contents

  1. How do I set up my computer?
  2. How is the code organized?
  3. How does Python talk to JavaScript?
  4. How does data get from GPU to the browser?
  5. How do I make changes?
  6. Patterns to know
  7. Anatomy of a widget file
  8. How do I add a new widget?
  9. How to make effective widgets
  10. Common mistakes to avoid
  11. Troubleshooting
  12. FAQ
  13. Glossary
  14. Testing
  15. Docs and publishing

How do I set up my computer? (5-10 minutes)

You need Python 3.11+, Node.js 20+, and the quantem core library. Choose one of the two setup options below.

Option A: conda (recommended if you already use conda/miniforge)

  1. Create and activate an environment:

    conda create -n widget python=3.12 nodejs=20 -y
    conda activate widget
  2. Install the core library:

    pip install quantem
  3. Clone and install the widget package:

    git clone https://github.com/bobleesj/quantem.widget.git
    cd quantem.widget
    npm install
    npm run build
    pip install -e .
  4. Verify it works:

    python -c "from quantem.widget import Show2D; print('OK')"

    You should see OK.

Option B: uv

  1. Install uv and Node.js 20+ (via nvm or your system package manager).

  2. Clone and install:

    git clone https://github.com/bobleesj/quantem.widget.git
    cd quantem.widget
    npm install
    npm run build
    uv sync
  3. Verify it works:

    uv run python -c "from quantem.widget import Show2D; print('OK')"

    You should see OK.

Note: If you're using uv, prefix Python commands with uv run (e.g., uv run jupyterlab), or activate the virtual environment with source .venv/bin/activate and use python directly.


How is the code organized?

Each widget is exactly one Python file + one TSX folder. No inheritance hierarchy - each widget is self-contained.

quantem.widget/
├── js/                          # TypeScript/React source
│   ├── align2d/index.tsx        # One folder per widget (each is self-contained)
│   ├── bin/index.tsx
│   ├── edit2d/index.tsx
│   ├── folder/index.tsx
│   ├── mark2d/index.tsx
│   ├── show1d/index.tsx
│   ├── show2d/index.tsx
│   ├── show3d/index.tsx
│   ├── show3dvolume/index.tsx
│   ├── show4d/index.tsx
│   ├── show4dstem/index.tsx
│   ├── showcomplex/index.tsx
│   ├── colormaps.ts             # ── Shared JS (used by all widgets) ──
│   ├── format.ts                #   Data extraction, number formatting
│   ├── histogram.ts             #   Histogram computation
│   ├── scalebar.ts              #   Scale bar, colorbar, figure export
│   ├── stats.ts                 #   Data statistics, log scale, clipping
│   ├── theme.ts                 #   Light/dark theme detection
│   ├── tool-parity.ts           #   Tool lock/hide logic
│   └── webgpu-fft.ts            #   GPU-accelerated FFT with CPU fallback
├── src/quantem/widget/          # Python source
│   ├── align2d.py               # One file per widget
│   ├── bin.py
│   ├── edit2d.py
│   ├── folder.py
│   ├── mark2d.py
│   ├── show1d.py
│   ├── show2d.py
│   ├── show3d.py
│   ├── show3dvolume.py
│   ├── show4d.py
│   ├── show4dstem.py
│   ├── showcomplex.py
│   ├── array_utils.py           # ── Shared Python (used by all widgets) ──
│   ├── json_state.py            #   State save/load envelope format
│   ├── tool_parity.py           #   Tool lock/hide validation
│   └── static/                  # Built JS bundles (generated by npm run build)
├── tests/                       # pytest unit tests + Playwright E2E tests
├── docs/                        # Sphinx documentation
│   └── examples/                #   Notebooks with saved widget state
├── notebooks/                   # Development/testing notebooks
├── package.json                 # esbuild config, npm scripts
└── pyproject.toml               # Python package config (hatch + hatch-jupyter-builder)

If you want to find where something lives, the rule is simple: Show4DSTEM's Python code is in src/quantem/widget/show4dstem.py, and its frontend code is in js/show4dstem/index.tsx.

Shared modules

Each widget is self-contained (one .py + one index.tsx), but a small set of utility modules are shared across all widgets. These exist because every widget needs them — they passed the bar of "used by 12+ widgets" before being extracted:

Python (in src/quantem/widget/):

  • array_utils.pyto_numpy() converts any input (NumPy, PyTorch, CuPy) to a NumPy array. Used by all 13 widgets.
  • json_state.pysave_state_file() and unwrap_state_payload() for the versioned JSON state envelope. Used by all 13 widgets.
  • tool_parity.py — validates disabled_tools / hidden_tools keys. Used by 12 widgets.

JavaScript (in js/):

  • theme.tsuseTheme() hook for automatic light/dark mode. Used by all 13 widgets.
  • colormaps.ts — colormap LUTs and applyColormap(). Used by all image widgets.
  • format.tsextractFloat32(), formatNumber(), downloadBlob(). Used by all widgets.
  • stats.tscomputeStats(), percentileClip(), applyLogScale(). Used by all image widgets.
  • scalebar.tsdrawScaleBarHiDPI(), exportFigure(). Used by all canvas widgets with calibration.
  • histogram.tscomputeHistogramFromBytes(). Used by all image widgets.
  • webgpu-fft.ts — GPU-accelerated 2D FFT with CPU fallback. Used by all FFT-capable widgets.
  • tool-parity.ts — tool lock/hide state management. Used by 12 widgets.

Do not create new shared modules unless the logic is already duplicated across many widgets. When in doubt, keep it in the widget file and extract later.

Widget catalog

Widget Description GPU
Show1D 1D trace viewer with multi-trace overlay, crosshair, legend No
Show2D 2D image viewer with gallery, ROI, line profiles, FFT, export No
Show3D 3D stack viewer with playback, ROI, FFT, export No
Show3DVolume Orthogonal slice viewer (XY, XZ, YZ) with FFT, playback No
Show4D General 4D dataset explorer with nav + signal panels, ROI masking Yes
Show4DSTEM 4D-STEM diffraction viewer with virtual imaging, ROI presets Yes
ShowComplex2D Complex-valued 2D viewer (amplitude/phase/HSV/real/imag) No
Mark2D Interactive 2D annotation (points, ROIs, profiles, snap-to-peak) No
Edit2D Interactive crop/pad/mask editor with brush tool No
Align2D Image alignment overlay with FFT-based auto-align No
Bin Detector binning with live preview and batch export Yes
Folder Folder browser for loading and viewing datasets No

How does Python talk to JavaScript?

anywidget bridges the two sides. You declare traitlets in Python, and they automatically sync to React state in the browser. Here's an overview - follow the numbered steps below, then use the diagram as a reference:

anywidget architecture: Python ↔ JavaScript bridge

Step 1: Python declares traitlets

Every trait marked .tag(sync=True) is visible to both Python and JavaScript:

class Show4DSTEM(anywidget.AnyWidget):
    pos_row = traitlets.Int(0).tag(sync=True)          # scalar - synced as number
    frame_bytes = traitlets.Bytes(b"").tag(sync=True)   # binary - synced as DataView

Step 2: JavaScript reads and writes traitlets

In the React component, useModelState() works like useState but the value is synced to the Python object:

const [posRow, setPosRow] = useModelState<number>("pos_row");
const [frameBytes] = useModelState<DataView>("frame_bytes");  // read-only

Calling setPosRow(5) in JavaScript sets pos_row = 5 on the Python object.

Step 3: Python reacts to changes

The .observe() method registers callbacks that fire when a trait changes:

self.observe(self._update_frame, names=["pos_row", "pos_col"])

When pos_row changes (from JS or Python), _update_frame() runs automatically.

Step 4: The full cycle

Here's what happens when you click a position on the scan image:

  1. JS: user clicks → setPosRow(5) fires
  2. Python: .observe() triggers _update_frame()
  3. Python: extracts the frame, sets self.frame_bytes = ...
  4. JS: receives new DataView, re-renders the canvas

This round-trip happens in milliseconds.

Trait types at a glance

Python type JS type Example use
traitlets.Int number Position, shape, indices
traitlets.Float number Calibration, ROI params
traitlets.Bool boolean Toggles (log scale, ROI active)
traitlets.Unicode string Mode selection, titles
traitlets.Bytes DataView Image data (raw float32)
traitlets.List Array Compound updates, statistics

How does data get from GPU to the browser?

The full dataset (potentially 16+ GB) stays on GPU or CPU as a PyTorch tensor. On each user interaction, Python extracts a single frame (~64 KB), serializes it as raw bytes, and sends it to the browser. The browser does colormapping and rendering entirely in JavaScript. Here's the flow:

Data flow: full dataset on GPU, single frame to browser

The key insight: a 16 GB dataset lives on GPU, but only 64 KB crosses the wire per click. The widget never sends the full dataset to the browser.

Where does the GPU work happen?

Show4DSTEM, Show4D, and Bin are the widgets that use PyTorch and GPU. They load data by converting to NumPy first, then moving to a PyTorch tensor on the best available device:

from quantem.core.config import validate_device

data_np = to_numpy(data)
device_str, _ = validate_device(None)  # Auto-detect: CUDA > MPS > CPU

# MPS can't handle tensors >INT_MAX elements; fall back to CPU
if data_np.size > 2**31 - 1 and device_str == "mps":
    device_str = "cpu"

self._device = torch.device(device_str)
self._data = torch.from_numpy(data_np.astype(np.float32)).to(self._device)

Device selection uses quantem.core.config.validate_device(), which reads the user's quantem configuration. The MPS fallback handles Apple Silicon's INT_MAX element limit — tensors over ~2 billion elements fall back to CPU (still fast on unified memory).

All other widgets (Show1D, Show2D, Show3D, Mark2D, etc.) are NumPy-only — no GPU.


How do I make changes?

Changing JavaScript (layout, rendering, UI)

  1. Start the watch mode:

    npm run dev

    This rebuilds JS bundles automatically when you save a .tsx file.

  2. Edit a file, for example js/show4dstem/index.tsx.

  3. Refresh the browser or re-run the notebook cell. No kernel restart needed.

Changing Python (data processing, observers)

  1. Make sure you installed in editable mode:

    pip install -e .

    You only need to do this once.

  2. Edit a file, for example src/quantem/widget/show4dstem.py.

  3. Restart the kernel and re-run cells. Python code is loaded at import time, so changes require a kernel restart.

A typical iteration cycle

  1. Open a notebook in notebooks/ (e.g., notebooks/show4dstem/test_show4dstem.ipynb)
  2. Edit js/show4dstem/index.tsx (JS) or src/quantem/widget/show4dstem.py (Python)
  3. JS changes: re-run the cell | Python changes: restart kernel, re-run
  4. Repeat

Patterns to know

Compound traits for batched updates

When the user drags an ROI on the diffraction panel, JS needs to send both the row and column at once. If you use two separate traits, Python fires two observers - one for each. Use a compound List trait instead:

roi_center = traitlets.List(traitlets.Float(), default_value=[0.0, 0.0]).tag(sync=True)
self.observe(self._on_roi_center_change, names=["roi_center"])

Now Python fires one observer per drag event, not two.

Binary data: always raw float32

Image data goes through the Bytes trait as raw float32 - no JSON, no base64:

# Python side: tensor → raw bytes
self.frame_bytes = frame.cpu().numpy().astype(np.float32).tobytes()
// JS side: DataView → Float32Array (uses extractFloat32 from format.ts)
import { extractFloat32 } from "../format";
const raw = extractFloat32(frameBytes);  // handles alignment + null check

Canvas rendering: layered architecture

Each widget stacks two or more <canvas> elements:

  1. Data canvas - colormapped image at native resolution (e.g., 128×128 px)
  2. UI canvas - crosshair, ROI shapes, scale bar, text overlays (drawn at device pixel ratio for crisp lines on HiDPI displays)

Some widgets add extra layers: Show3D has an overlay canvas + lens canvas, Show4DSTEM has separate data/overlay/UI sets per panel (DP, virtual image, FFT). The key principle: data pixels live on the data canvas at image resolution, all text and vector overlays live on HiDPI canvases so they stay sharp at any zoom.

Coordinate convention: (row, col), never (x, y)

All user-facing coordinates use (row, col). This matches electron microscopy convention and NumPy array indexing (array[row, col]). The first value is always the row (vertical, top-to-bottom), the second is the column (horizontal, left-to-right).

This applies everywhere:

  • Python API: trait names use roi_row/roi_col, not roi_x/roi_y. Point dicts use {"row": r, "col": c}. Tuple inputs are (row, col).
  • JS/TS types: Point = { row, col }, ROI = { row, col }.
  • Display: coordinates shown as (row, col) in stats bars, tooltips, and readouts.

Internal drawing code (canvas pixel positions, DOM events, pan/zoom state) can use x/y since those are screen coordinates, not image coordinates. But anything exposed to the user - traits, dicts synced to Python, hover readouts, summary() output - must use (row, col).

# Right
pos_row = traitlets.Int(0).tag(sync=True)
pos_col = traitlets.Int(0).tag(sync=True)

# Wrong - don't use x/y for image coordinates
pos_x = traitlets.Int(0).tag(sync=True)
pos_y = traitlets.Int(0).tag(sync=True)

Array compatibility

All widgets accept NumPy arrays, PyTorch tensors, CuPy arrays, and any object compatible with np.asarray(). Conversion is handled by array_utils.to_numpy() via duck typing - no hard dependency on PyTorch or CuPy.

Widgets also auto-extract metadata from quantem Dataset objects (title, pixel size, units) using hasattr checks. Explicit parameters always override auto-extracted values:

# All of these work
w = Show2D(numpy_array)
w = Show2D(torch_tensor)
w = Show2D(cupy_array)
w = Show2D(quantem_dataset)            # auto-extracts title, pixel_size
w = Show2D(quantem_dataset, title="Override")  # explicit wins

Replacing data with set_image()

Every widget implements set_image() (or set_images() for Align2D) to swap in new data while preserving display settings. This avoids recreating the widget when data changes:

w = Show2D(image_a, cmap="viridis", log_scale=True)

# Later - swap data, keep cmap and log_scale
w.set_image(image_b)

Each widget defines what it preserves (cmap, log_scale, ROI settings) and what it resets (slice index, computed stats). See the widget docstring for details.

Scale bar: HiDPI rendering

All canvas-based widgets draw a scale bar on the UI canvas (the top layer), not the data canvas. This keeps text crisp at any zoom level. The pattern:

  • Render at device pixel ratio (width={cssW * DPR}, style={{ width: cssW }})
  • Fixed CSS sizes: 60 px bar, 5 px thick, 16 px font, 12 px margin
  • Always white with drop shadow for contrast on any colormap
  • Physical length computed from pixel size and zoom level, rounded to a nice value
  • Unit conversion: angstroms to nm at 10+ angstroms, px for uncalibrated data

Scale bar rendering is shared via drawScaleBarHiDPI() in js/scalebar.ts.


Anatomy of a widget file

Every widget follows the same structure. Here's how Show4DSTEM is organized — the most complex widget (~4200 lines Python, ~4100 lines TSX). Simpler widgets like Show2D have the same sections, just shorter.

Python file (src/quantem/widget/show4dstem.py)

Read top-to-bottom, the file has 7 sections in this order:

┌─────────────────────────────────────────────────┐
│  1. Imports + constants                         │
│  2. Trait declarations (synced to JS)           │
│  3. __init__ (data loading, setup)              │
│  4. Observers (react to trait changes)          │
│  5. Public API (set_image, ROI presets, etc.)   │
│  6. State protocol (state_dict, save, summary)  │
│  7. Internal methods (_get_frame, _render, etc.)│
└─────────────────────────────────────────────────┘

1. Imports + constants — Standard imports, anywidget, traitlets, numpy, torch, shared modules (to_numpy, save_state_file, tool parity). Constants like DEFAULT_BF_RATIO, MIN_LOG_VALUE.

2. Trait declarations — Every synced property in one block at the top of the class. Grouped by purpose:

class Show4DSTEM(anywidget.AnyWidget):
    # Position & shape
    pos_row = traitlets.Int(0).tag(sync=True)
    pos_col = traitlets.Int(0).tag(sync=True)
    shape_rows = traitlets.Int(0).tag(sync=True)
    frame_bytes = traitlets.Bytes(b"").tag(sync=True)

    # Display settings
    dp_colormap = traitlets.Unicode("inferno").tag(sync=True)
    dp_scale_mode = traitlets.Unicode("linear").tag(sync=True)

    # ROI
    roi_mode = traitlets.Unicode("off").tag(sync=True)
    roi_center_row = traitlets.Float(0.0).tag(sync=True)
    ...

This is the single source of truth for what Python and JavaScript share. If a trait isn't here, JS can't see it. If it lacks .tag(sync=True), JS can't see it.

3. __init__ — The constructor. Always follows this pattern:

def __init__(self, data, ..., state=None, **kwargs):
    super().__init__(**kwargs)          # 1. Initialize anywidget
    self.widget_version = resolve_widget_version()

    data_np = to_numpy(data)           # 2. Convert input to NumPy
    # ... shape validation ...          # 3. Validate dimensions

    self._device = torch.device(...)   # 4. Move to GPU (if applicable)
    self._data = torch.from_numpy(...).to(self._device)

    self.shape_rows = ...              # 5. Set traits from data
    self.frame_bytes = ...             # 6. Send first frame to JS

    self.observe(self._update_frame,   # 7. Register observers
                 names=["pos_row", "pos_col"])

    if state is not None:              # 8. Restore saved state (last!)
        self.load_state_dict(state)

Key rule: super().__init__(**kwargs) must be first. State restoration must be last (after all traits are set).

4. Observers — Functions triggered by trait changes. Registered in __init__ via self.observe():

self.observe(self._update_frame, names=["pos_row", "pos_col"])
self.observe(self._update_virtual_image, names=["roi_mode", "roi_center"])

When JS changes pos_row (user clicks), _update_frame() runs in Python, extracts the frame, and sets self.frame_bytes — which JS then renders.

5. Public API — Methods users call from Python: set_image(), roi_circle(), set_profile(), set_path(), play(), pause(), etc. These return self for chaining:

w = Show4DSTEM(data)
w.roi_circle().set_path(points).play()

6. State protocolstate_dict(), save(), load_state_dict(), summary(), __repr__(). Same pattern in every widget. state_dict includes display settings and ROI, excludes data and computed stats.

7. Internal methods — Prefixed with _. Data extraction (_get_frame), rendering helpers (_render_dp_rgb), export logic, computation. Not part of the public API.

TypeScript file (js/show4dstem/index.tsx)

Also read top-to-bottom, the file has 7 sections:

┌─────────────────────────────────────────────────┐
│  1. Imports + constants + styles                │
│  2. Utility functions + helper components       │
│  3. Model state hooks (useModelState)           │
│  4. Local state + refs                          │
│  5. Effects (useEffect) — rendering + setup     │
│  6. Event handlers (mouse, keyboard, export)    │
│  7. JSX return (the actual UI layout)           │
└─────────────────────────────────────────────────┘

1. Imports + constants + styles — React, MUI components, shared modules. Then constants: zoom limits, font sizes, spacing, colors, button styles. These are defined outside the component function so they don't recreate on each render:

const MIN_ZOOM = 0.5;
const MAX_ZOOM = 40;
const CANVAS_SIZE = 360;
const SPACING = { xs: "2px", sm: "4px", md: "8px" };
const typography = { label: { fontSize: 11 }, stat: { fontSize: 10, fontFamily: "monospace" } };
const compactButton = { fontSize: 10, minWidth: 0, px: "6px", py: "2px" };

2. Utility functions + helper components — Small components and drawing functions defined outside the main component: Histogram (sparkline with interactive range), InfoTooltip (themed tooltip), KeyboardShortcuts (shortcut table), and canvas drawing helpers (drawRoiOverlayHiDPI, drawViPositionMarker).

3. Model state hooks — Inside the component function, the first block reads synced traits from Python:

function Show4DSTEM({ model }: { model: AnyModel }) {
  const [posRow, setPosRow] = useModelState<number>("pos_row");
  const [frameBytes] = useModelState<DataView>("frame_bytes");
  const [dpColormap, setDpColormap] = useModelState<string>("dp_colormap");
  // ... 40+ more useModelState calls

This mirrors the Python trait declarations. Same names, same types.

4. Local state + refs — UI-only state that doesn't sync to Python: drag tracking, cursor position, canvas dimensions, animation frames. Plus refs to canvas DOM elements:

const [canvasSize, setCanvasSize] = useState(CANVAS_SIZE);
const [cursorInfo, setCursorInfo] = useState<CursorInfo | null>(null);
const dpCanvasRef = React.useRef<HTMLCanvasElement>(null);
const dragRef = React.useRef({ startX: 0, startY: 0, wasDrag: false });

5. Effects (useEffect) — The rendering engine. Each effect watches specific dependencies and re-runs when they change:

// When frame data or colormap changes → re-render the diffraction panel
React.useEffect(() => {
  const raw = extractFloat32(frameBytes);
  if (!raw) return;
  const scaled = applyScale(raw, scaleMode);
  const rgba = applyColormap(scaled, lut, vmin, vmax);
  // ... draw to canvas
}, [frameBytes, dpColormap, dpScaleMode, dpVminPct, dpVmaxPct, zoom, panX, panY]);

Effects handle: parsing binary data, applying colormaps, FFT computation, drawing scale bars/ROIs/crosshairs, setting up scroll prevention, animation timers.

6. Event handlers — Mouse (click, drag, zoom), keyboard (arrow keys, R for reset), export (PNG, figure, GIF). Each panel has its own handlers:

const handleDpMouseDown = (e: React.MouseEvent) => { ... };
const handleDpMouseMove = (e: React.MouseEvent) => { ... };
const handleViMouseDown = (e: React.MouseEvent) => { ... };
const handleKeyDown = (e: React.KeyboardEvent) => {
  if (e.key === "ArrowRight") setPosCol(Math.min(posCol + 1, shapeCols - 1));
  if (e.key === "r" || e.key === "R") resetZoom();
};

7. JSX return — The actual UI layout. Builds the widget from top to bottom:

return (
  <Box sx={{ ...container, bgcolor: colors.bg }}>
    {/* Header: title + info tooltip + buttons */}
    <Stack direction="row">...</Stack>

    {/* Panels: DP | VI | FFT (side by side) */}
    <Stack direction="row">
      {/* DP panel: canvas stack + stats + controls */}
      <Box>
        <canvas ref={dpCanvasRef} />      {/* data layer */}
        <canvas ref={dpOverlayRef} />     {/* overlay layer */}
        <canvas ref={dpUiRef} />          {/* HiDPI UI layer */}
        <Box>Mean: ... Min: ... Max: ...</Box>  {/* stats bar */}
        <Box>                             {/* controls */}
          <Select>Colormap</Select>
          <Switch>Log</Switch>
          <Histogram />
        </Box>
      </Box>

      {/* VI panel: same structure */}
      {/* FFT panel: conditional */}
    </Stack>

    {/* Playback controls (frame slider, path slider) */}
  </Box>
);

The canvas layers are stacked with position: absolute — data canvas at image resolution, overlay for interactive drawing, UI canvas at HiDPI resolution for crisp text.

How the sections connect

Here's the full cycle when a user clicks on the virtual image panel:

JS: handleViMouseDown          → setPosRow(5), setPosCol(10)
                                         ↓
Python: self.observe() fires   → _update_frame(pos_row=5, pos_col=10)
                                         ↓
Python: _get_frame()           → self.frame_bytes = tensor[5, 10].tobytes()
                                         ↓
JS: useEffect [frameBytes]     → extractFloat32 → applyColormap → drawImage
                                         ↓
Browser: canvas shows new diffraction pattern

Every widget follows this exact cycle. The specifics (which traits, which canvas, which colormap) change, but the structure is always: JS event → trait update → Python observer → new bytes → JS effect → canvas render.


How do I add a new widget?

Say you want to add a widget called ShowPDF. Here's the approach that works best:

Start from the UI, not the data

Before writing any code, draw what you want the widget to look like. Sketch it on paper, in Figma, or just describe it in words. Figure out:

  • What does the user see? (canvases, sliders, buttons, panels)
  • What can the user click or drag?
  • What should happen when they do?

Once you know the UI, pick the existing widget that's closest to what you want and copy it wholesale. Show1D (~450 lines Python, ~1200 lines TSX) is the simplest canvas widget. Show2D (~1100 lines Python, ~3000 lines TSX) is a good mid-complexity starting point. Show4DSTEM is the most complex (~4200 lines Python, ~4100 lines TSX). Start simple.

Tip: Copy the model widget's UI code into your favorite LLM (e.g., Claude or ChatGPT), describe what you want to change, and let it modify the layout, controls, and rendering for you. This is much faster than writing TSX from scratch, especially if you're not fluent in React.

Files to create

  1. src/quantem/widget/showpdf.py - the Python backend:

    • Extend anywidget.AnyWidget
    • Set _esm and _css pointing to static/showpdf.js and static/showpdf.css
    • Define traitlets with .tag(sync=True)
  2. js/showpdf/index.tsx - the JavaScript frontend:

    • Use useModelState() to read/write traitlets
    • Render to <canvas> elements

Files to edit

  1. package.json - add the entry point to the build command:

    js/showpdf/index.tsx
    
  2. pyproject.toml - add the output to ensured-targets:

    "src/quantem/widget/static/showpdf.js"
    
  3. src/quantem/widget/__init__.py - export ShowPDF.

Files to add for testing and docs

  1. tests/test_widget_showpdf.py - basic trait and instantiation tests.

  2. docs/widgets/showpdf.rst and a notebook in docs/examples/showpdf/.

Verify it works

npm run build                    # bundle showpdf.js
pip install -e .                 # reinstall with new widget
python -c "from quantem.widget import ShowPDF; print('OK')"
pytest tests/test_widget_showpdf.py

How to make effective widgets

Our users are scientists. They care about their data, not our UI framework. A good widget disappears - the scientist interacts with the data, not with buttons and panels. Here's what that means in practice:

Think in terms of "what does the scientist need to see right now?" A 4D-STEM dataset is 4 GB. The scientist clicking a scan position needs to see one 128×128 diffraction pattern. That's 64 KB. Send 64 KB, not 4 GB. The Bytes trait handles raw binary - no JSON, no base64, no serialization overhead. Every widget should have this property: the data on the wire is proportional to what's on screen, not what's in memory.

Python does computation. JavaScript does rendering. Don't mix them. Python owns the data: slicing tensors, computing virtual images, running GPU kernels. JavaScript owns what the user sees: colormaps, canvas drawing, zoom/pan, crosshairs, tooltips, and even GPU-accelerated FFT (via WebGPU). If you're building HTML strings in Python, stop. The boundary is the Bytes trait.

Keep interactive feedback on the JS side. Mouse tracking, hover coordinates, zoom/pan - all of these run at 60fps in the browser with zero round-trips. Only cross to Python when the scientist does something that requires recomputation (clicking a new scan position, changing an ROI). The round-trip is fast (~ms), but the browser is faster (~16ms frames).

Test with real data from day one. A 10×10×10×10 test array won't reveal the problems a 256×256×128×128 dataset will. Load a real .h5 file, click around, drag the ROI. Performance issues, memory problems, and edge cases all surface with real data. Scientists will use real data - so should you.


Common mistakes to avoid

Don't over-engineer the file structure. Each widget is one Python file and one TSX folder. That's it. show4dstem.py is ~4200 lines and that's fine - it's one self-contained unit with traitlets, observers, and computation in one place. The temptation is to split it into show4dstem_traits.py, show4dstem_compute.py, show4dstem_observers.py. Don't. Splitting a single widget across files makes it harder to understand, not easier. One file means you can read top-to-bottom and see everything the widget does. It also means you can feed the entire widget to an LLM in one shot - no chasing imports across five files to debug a single observer callback.

Don't create abstract base classes for widgets. There's no BaseWidget or GPUWidget parent class. Each widget stands alone. Show4DSTEM and Show4D both use to_numpy() from array_utils.py - that's shared via a utility function, not inheritance. Minimize sharing between widgets unless you know the function will be used by many - potentially all - widgets. to_numpy() qualifies because every widget needs it. A one-off helper used by two widgets does not. When in doubt, duplicate. Extract only when the pattern is proven across dozens of call sites, like FFT or colormap application. Get the widget working first - cleanup and refactoring can be done easily later with an LLM.

Don't over-split the JavaScript side either. Each widget has one entry point: js/show4dstem/index.tsx. Shared JS utilities (colormaps.ts, scalebar.ts, theme.ts, stats.ts) exist at the js/ level because multiple widgets import them. These shared modules exist because every widget needs colormaps, scale bars, and theme detection - they're universal concerns, not premature abstractions. Don't create a new shared module unless multiple widgets already use the same logic and future widgets will too. If you want to add a new shared module, discuss with @bobleesj via GitHub issue or Slack first.

Don't send the full dataset to the browser. If you assign the entire tensor to a Bytes trait, the browser tab will freeze or crash. Only send the slice the browser needs for the current view.

Don't add unnecessary dependencies. Before implementing a computation (CoM, FFT, alignment, masking), check if it already exists in quantem core. The widget layer should call existing functions, not reimplement them.


Troubleshooting

Trait exists in Python but JavaScript reads undefined

You forgot .tag(sync=True). Without it, the traitlet is Python-only:

# Wrong - JS can't see this
pos_row = traitlets.Int(0)

# Right
pos_row = traitlets.Int(0).tag(sync=True)
Python changes don't take effect

Restart the kernel. JavaScript changes take effect on cell re-run, but Python code is loaded at import time. If your edit doesn't seem to work, restart the kernel before debugging further.

Trait isn't syncing between Python and JS

Check the spelling. The string in useModelState("pos_row") must exactly match the Python trait name pos_row. A typo means JS silently reads undefined. Check both sides.


FAQ

How does the widget detect light vs. dark mode?

Every widget calls useTheme() from theme.ts. It auto-detects JupyterLab, VS Code, Colab, Classic Jupyter, and OS-level preferences - no configuration needed. It returns a colors object and a themeInfo with the detected environment and theme:

const { themeInfo, colors } = useTheme();
// themeInfo.theme is "light" or "dark"
// themeInfo.environment is "jupyterlab", "vscode", "colab", etc.
How do I use colors in my widget?

Use the colors object from useTheme() for all UI elements - backgrounds, text, borders, controls. This keeps every widget consistent and theme-aware:

const { colors } = useTheme();

// Backgrounds and text
<Box sx={{ bgcolor: colors.bg, color: colors.text }}>

// Borders and controls
<Box sx={{ border: `1px solid ${colors.border}`, bgcolor: colors.controlBg }}>

// Muted text and accents
<Typography sx={{ color: colors.textMuted }}>
  Value: <span style={{ color: colors.accent }}>{value}</span>
</Typography>

The available colors are: bg, bgAlt, text, textMuted, border, controlBg, and accent. These adapt automatically between light and dark mode.

What if I need a specific color palette (not from the theme)?

That's fine - use it directly. Scientific visualizations often need domain-specific colors (e.g., red/blue for strain, specific colormaps for diffraction). Use theme colors for the UI chrome (panels, labels, borders) and your own palette for the data visualization itself.

How do colormaps work?

colormaps.ts provides LUTs (lookup tables) for standard scientific colormaps: inferno, viridis, plasma, magma, hot, gray, and hsv. Use applyColormap() to map float32 data to RGBA pixels:

import { COLORMAPS, applyColormap } from "../colormaps";

const lut = COLORMAPS["inferno"];
applyColormap(floatData, rgbaBuffer, lut, vmin, vmax);

The colormap is applied entirely in JavaScript - Python sends raw float32 bytes, and JS handles the visual mapping.


Glossary

You'll see these terms throughout the widget code. Organized by where you'll encounter them.

Python terms

traitlets

The library that defines typed attributes on Python objects. Every widget property that syncs to JavaScript is a traitlet. Think of them as "Python properties with type checking and change notifications."

# Declares an integer trait with default 0, synced to JS
pos_row = traitlets.Int(0).tag(sync=True)

Common types: Int, Float, Bool, Unicode (string), Bytes (binary), List, Dict.

.tag(sync=True)

Marks a traitlet for synchronization with JavaScript. Without this tag, the trait exists only in Python — JavaScript can't see it. Every trait that appears in the UI must have .tag(sync=True).

.observe(callback, names=[...])

Registers a function that runs whenever the named traits change. This is how Python reacts to user interactions — JS updates a trait, and .observe() fires the callback:

self.observe(self._update_frame, names=["pos_row", "pos_col"])

The callback receives a change dict with old and new values, but most widgets ignore it and just read the current trait values directly.

anywidget.AnyWidget

The base class every widget extends. It handles the Python↔JS bridge: trait synchronization, binary transport, and widget lifecycle. You subclass it, define traitlets, and point _esm / _css to your JS/CSS files.

@property

A Python decorator that turns a method into a read-only attribute. Widgets use this for computed values:

@property
def profile_values(self):
    return np.frombuffer(self._profile_bytes, dtype=np.float32)

Now widget.profile_values works like an attribute, not a method call.

Self return type

A type hint meaning "returns the same object." Used for method chaining:

def set_image(self, data) -> Self:
    # ... update data ...
    return self

# Enables: widget.set_image(data).set_profile(p0, p1)

Imported from typing (Python 3.11+).

super().__init__(**kwargs)

Calls the parent class constructor. Every widget's __init__ must call this to initialize the anywidget machinery. Always pass **kwargs through so anywidget can handle its internal arguments.

.tobytes()

Converts a NumPy array to raw bytes. This is how image data goes from Python to JavaScript — as a flat binary buffer with no metadata:

self.frame_bytes = frame.astype(np.float32).tobytes()

On the JS side, this arrives as a DataView.

hasattr(obj, "attr") (duck typing)

Checks if an object has a given attribute without caring about its type. Widgets use this to auto-detect quantem Dataset objects without importing quantem:

if hasattr(data, "array"):
    # It's a Dataset — extract the array and metadata
    title = getattr(data, "name", "")

This avoids a hard dependency on the quantem package.

JavaScript / TypeScript terms

JSX

HTML-like syntax inside JavaScript. React uses it to describe UI. The <Box>, <canvas>, <Typography> tags you see in .tsx files are JSX — they compile to JavaScript function calls:

// This JSX:
<Box sx={{ color: "red" }}>Hello</Box>

// Compiles to:
React.createElement(Box, { sx: { color: "red" } }, "Hello")

You write JSX like HTML but with {} for JavaScript expressions.

Hooks (useState, useEffect, useRef, useMemo, useCallback)

Functions that manage state and side effects in React components. Every widget uses these extensively:

  • useState(initial) — declares a state variable. Returns [value, setValue]. When you call setValue(newValue), the component re-renders.

  • useEffect(() => { ... }, [deps]) — runs code when dependencies change. This is where widgets set up canvas rendering, event listeners, and data processing. The dependency array [deps] controls when it re-runs.

  • useRef(initial) — creates a persistent reference that survives re-renders. Used for canvas elements (canvasRef.current), mutable state that shouldn't trigger re-renders (drag state, animation frames), and previous values.

  • useMemo(() => value, [deps]) — caches a computed value. Recalculates only when dependencies change. Used for expensive operations like parsing float32 data from DataView.

  • useCallback(() => fn, [deps]) — caches a function reference. Similar to useMemo but for functions. Used for event handlers to prevent unnecessary re-renders.

useModelState

The anywidget hook that bridges React to Python traitlets. Works like useState but the value lives in the Python object:

const [cmap, setCmap] = useModelState<string>("cmap");
// cmap reads the Python trait value
// setCmap("viridis") updates the Python trait

The string must exactly match the Python trait name. A typo silently returns undefined.

DataView and Float32Array

JavaScript types for working with binary data.

  • DataView — a generic wrapper around binary data. This is what Bytes traitlets arrive as in JavaScript. You can't use it directly for math.

  • Float32Array — a typed array of 32-bit floats. This is what you need for image data. Convert from DataView using extractFloat32() from format.ts:

const raw = extractFloat32(frameBytes);  // DataView → Float32Array
// Now raw[i] gives you pixel values as numbers
Canvas API (getContext, drawImage, clearRect)

The browser's 2D drawing API. Every widget renders to <canvas> elements using this:

const ctx = canvas.getContext("2d");
ctx.clearRect(0, 0, w, h);           // clear the canvas
ctx.drawImage(offscreen, 0, 0, w, h); // draw an image onto the canvas
ctx.beginPath();                       // start drawing a shape
ctx.arc(x, y, radius, 0, 2 * Math.PI); // circle
ctx.fill();                            // fill the shape
ctx.save();    // save current transform
ctx.restore(); // restore it later

Offscreen canvas — a canvas that isn't visible on screen. Widgets render the colormapped image onto an offscreen canvas first, then drawImage() it onto the visible canvas with zoom/pan transforms applied.

DPR (device pixel ratio)

window.devicePixelRatio — how many physical pixels per CSS pixel. On a Retina MacBook, DPR = 2 (each CSS pixel is 2×2 physical pixels).

Widgets use this for the UI canvas to keep text and lines crisp:

const DPR = window.devicePixelRatio || 1;
// Canvas element: width={cssW * DPR}, height={cssH * DPR}
// CSS style: width={cssW}, height={cssH}
// Drawing: ctx.scale(DPR, DPR) — then draw in CSS pixel units

Without DPR scaling, text and lines look blurry on HiDPI displays.

ref / .current

A ref is a persistent reference to a DOM element or mutable value. Created with useRef(), accessed via .current:

const canvasRef = React.useRef<HTMLCanvasElement>(null);

// In JSX — attaches the ref to the canvas element:
<canvas ref={canvasRef} />

// In code — access the actual DOM element:
const ctx = canvasRef.current?.getContext("2d");

Refs don't trigger re-renders when they change. This makes them ideal for canvas elements, animation state, and drag tracking.

getBoundingClientRect()

Returns the position and size of a DOM element on screen. Widgets use this to convert mouse positions to canvas coordinates:

const rect = canvas.getBoundingClientRect();
const mouseX = e.clientX - rect.left;  // mouse X relative to canvas
const mouseY = e.clientY - rect.top;   // mouse Y relative to canvas
ImageData

A pixel buffer that the canvas can display. Contains RGBA values (4 bytes per pixel). This is what the colormap pipeline produces:

const imageData = new ImageData(rgbaBuffer, width, height);
ctx.putImageData(imageData, 0, 0);

However, most widgets use drawImage() from an offscreen canvas instead of putImageData(), because drawImage() supports transforms (zoom/pan).

MUI (Material UI) components

Box, Stack, Typography

MUI layout components. Think of them as styled <div> elements:

  • Box — a <div> with MUI's sx styling. The most common layout element.
  • Stack — a Box with flexbox layout. direction="row" for horizontal, direction="column" for vertical.
  • Typography — a text element with consistent styling.
<Stack direction="row" spacing={1} alignItems="center">
  <Typography sx={{ fontSize: 10 }}>Scale:</Typography>
  <Select size="small" value={scale}>...</Select>
</Stack>
sx prop

MUI's inline styling system. Like CSS but written as a JavaScript object:

// CSS:        font-size: 10px; margin-bottom: 4px; background-color: #fff;
// sx:
<Box sx={{ fontSize: 10, mb: "4px", bgcolor: "#fff" }}>

Shorthand properties: mb (margin-bottom), mt (margin-top), p (padding), bgcolor (background-color). Full CSS properties also work.

Switch, Select, MenuItem, Button, Slider

MUI form controls:

  • Switch — toggle on/off (used for FFT, log scale, auto contrast)
  • Select + MenuItem — dropdown (used for colormap, scale mode)
  • Button — clickable button (used for Reset, Export, Copy)
  • Slider — draggable range (used for contrast, frame scrubbing)

All styled via sx to match the widget's compact layout (fontSize 10, small switches, no border radius).

Build tools

esbuild

The JavaScript bundler. Takes .tsx files from js/ and produces .js bundles in src/quantem/widget/static/. Configured in package.json. Very fast (~130ms full build).

  • npm run build — one-shot build
  • npm run dev — watch mode (rebuilds on file save)
anywidget HMR

Hot Module Replacement — when npm run dev is running and ANYWIDGET_HMR=1 is set, JavaScript changes appear in the notebook without restarting the kernel or re-running cells. The widget's JS code reloads in place.


Testing

Run unit tests (fast, no browser needed):

pytest tests/ --ignore=tests/test_e2e_smoke.py

Run tests for one widget:

pytest tests/test_widget_show4dstem.py

Run end-to-end smoke tests (requires Playwright + JupyterLab, ~4 min):

pytest tests/test_e2e_smoke.py -v

What the test files do:

  • test_widget_*.py - unit tests (trait initialization, shape validation, data loading, state persistence, ROI, display traits)
  • capture_*.py - screenshot capture scripts for visual regression (run manually in JupyterLab to generate reference images)
  • test_array_utils.py - tests for to_numpy() array conversion
  • test_e2e_smoke.py - end-to-end tests using Playwright: renders each widget in a real browser, captures screenshots (light + dark theme), and tests interactions (colormap changes, FFT toggles, log scale, ROI). Takes ~4 minutes.

Docs and publishing

Building docs locally

Notebooks in docs/examples/ are rendered by Sphinx with nbsphinx. Widget state is pre-saved in the notebook (File → Save in JupyterLab), so docs pages show interactive widgets without a running kernel:

pip install -e ".[docs]"
sphinx-build docs docs/_build/html

Publishing to TestPyPI

  1. Bump the version in pyproject.toml
  2. Commit, tag, and push:
    git tag v0.0.6
    git push origin main && git push origin v0.0.6

GitHub Actions builds JS, packages the wheel, and uploads to TestPyPI. See README.md for full CI/CD details including trusted publisher setup.