# ShowComplex2D — Quick Demo Interactive viewer for complex-valued data from ptychography, SSB, holography, and exit wave reconstruction. Five display modes: amplitude, phase, HSV, real, imaginary.

# Install in Google Colab
try:
    import google.colab
    !pip install -q -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ quantem-widget
except ImportError:
    pass  # Not in Colab, skip

try:
    %load_ext autoreload
    %autoreload 2
    %env ANYWIDGET_HMR=1
except Exception:
    pass  # autoreload unavailable (Colab Python 3.12+)

env: ANYWIDGET_HMR=1

import torch
import numpy as np
import quantem.widget
from quantem.widget import ShowComplex2D
device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
def make_ptycho_object(size=256):
    """Simulate a ptychographic reconstruction of a SrTiO3 [001] zone axis.
    Object transmission function: amplitude ≈ 1 (thin specimen, weak absorption),
    phase = projected atomic potential with realistic column positions.
    """
    y = torch.linspace(0, 1, size, device=device).unsqueeze(1).expand(size, size)
    x = torch.linspace(0, 1, size, device=device).unsqueeze(0).expand(size, size)
    # SrTiO3-like perovskite: Sr on corners, Ti on body center, O on face centers
    # Unit cell repeats ~8 times across the field of view
    n_cells = 8
    phase = torch.zeros(size, size, device=device)
    # Atomic column potential: sum of Gaussians at lattice sites
    for i in range(n_cells + 1):
        for j in range(n_cells + 1):
            cx, cy = i / n_cells, j / n_cells
            r2 = (x - cx) ** 2 + (y - cy) ** 2
            # Sr columns (corners): strong phase, Z=38
            phase += 0.8 * torch.exp(-r2 / (2 * (0.012) ** 2))
            # Ti columns (body center): medium phase, Z=22
            tcx, tcy = (i + 0.5) / n_cells, (j + 0.5) / n_cells
            r2_ti = (x - tcx) ** 2 + (y - tcy) ** 2
            phase += 0.5 * torch.exp(-r2_ti / (2 * (0.010) ** 2))
            # O columns (face centers): weak phase, Z=8
            for ox, oy in [(cx + 0.5 / n_cells, cy), (cx, cy + 0.5 / n_cells)]:
                if ox <= 1.0 and oy <= 1.0:
                    r2_o = (x - ox) ** 2 + (y - oy) ** 2
                    phase += 0.15 * torch.exp(-r2_o / (2 * (0.008) ** 2))
    # Amplitude: near 1 with slight absorption at heavy columns
    amp = 1.0 - 0.08 * phase / phase.max()
    # Add subtle scan noise (realistic for ptychography)
    noise = 0.02 * torch.randn(size, size, device=device)
    phase = phase + noise
    obj = amp * torch.exp(1j * phase)
    return obj.cpu().numpy()
obj = make_ptycho_object(256)
print(f"Ptycho object: shape={obj.shape}, dtype={obj.dtype}")
print(f"quantem.widget {quantem.widget.__version__}")

Using device: mps
Ptycho object: shape=(256, 256), dtype=complex64
quantem.widget 0.4.0a3

# Phase view: projected atomic potential (most common ptychography visualization)
ShowComplex2D(obj, title="SrTiO₃ Ptycho — Phase", display_mode="phase", pixel_size_angstrom=0.5)

# Amplitude view: object transmission (close to 1 for thin specimen)
ShowComplex2D(obj, title="SrTiO₃ Ptycho — Amplitude", display_mode="amplitude", pixel_size_angstrom=0.5)

# HSV view: phase→hue, amplitude→brightness (simultaneous visualization)
ShowComplex2D(obj, title="SrTiO₃ Ptycho — HSV", display_mode="hsv", pixel_size_angstrom=0.5)

Inspect Widget State

w = ShowComplex2D(obj, title="SrTiO₃ Ptychographic Reconstruction", display_mode="phase", pixel_size_angstrom=0.5)
w.summary()
w

SrTiO₃ Ptychographic Reconstruction
════════════════════════════════
Image:    256×256 (complex)
Amp:      min=0.92  max=1  mean=0.9925
Phase:    min=-0.08244  max=0.8437  mean=0.0754
Display:  phase | hsv (cyclic) | manual | linear