"""
Quantum STEM (Scanning Transmission Electron Microscopy)
=========================================================
Simulates STEM images using TRUE quantum electron probe propagation.
Each probe position runs a quantum circuit (DiagonalGate + QFT).
Detectors integrate scattered intensity over defined angular ranges:
HAADF — High-Angle Annular Dark Field (Z-contrast)
ADF — Annular Dark Field
ABF — Annular Bright Field
BF — Bright Field
iDPC — integrated Differential Phase Contrast
Physical steps per probe position:
1. Coherent focused probe formed in k-space with CTF.
2. Phase grating applied via quantum DiagonalGate circuit.
3. Free-space propagation in k-space via diagonal phase (Fresnel).
4. Detector masks applied → signal readout.
For large grids (n_qubits > MAX_SV_QUBITS) or large scan arrays,
a classical numpy fallback is used automatically.
References
----------
- Kirkland (2010). Advanced Computing in Electron Microscopy.
- Nellist & Pennycook (1999). Incoherent imaging. Adv. Imaging Elec. Phys. 113.
- Ophus (2023). 4D-STEM. arXiv:2301.00345.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
import numpy as np
from qiskit import QuantumCircuit
from qiskit.circuit.library import DiagonalGate
from qiskit.quantum_info import Statevector
from quscope.quantum_ctem.quantum_ctem_circuit import (
relativistic_wavelength,
interaction_constant,
)
MAX_SV_QUBITS = 14 # Statevector feasible up to this qubit count
# ── Detector definitions ──────────────────────────────────────────────────────
[docs]
@dataclass
class STEMDetectors:
"""
Angular detector definitions for STEM.
All angles in mrad. Pass to run_stem().
"""
haadf_inner: float = 60.0 # HAADF inner angle (mrad)
haadf_outer: float = 200.0 # HAADF outer angle (mrad)
adf_inner: float = 25.0 # ADF inner angle
adf_outer: float = 60.0 # ADF outer angle
abf_inner: float = 10.0 # ABF inner angle
abf_outer: float = 25.0 # ABF outer angle
bf_outer: float = 10.0 # BF half-angle
[docs]
def masks(
self,
N: int,
pixel_size: float,
wavelength: float,
) -> Dict[str, np.ndarray]:
"""Return boolean k-space masks for each detector."""
freq = np.fft.fftshift(np.fft.fftfreq(N, d=pixel_size))
KX, KY = np.meshgrid(freq, freq, indexing="ij")
K = np.sqrt(KX ** 2 + KY ** 2)
K_mrad = K * wavelength * 1e3 # k in mrad
def annular(inner, outer):
return (K_mrad >= inner) & (K_mrad < outer)
return {
"HAADF": annular(self.haadf_inner, self.haadf_outer),
"ADF": annular(self.adf_inner, self.adf_outer),
"ABF": annular(self.abf_inner, self.abf_outer),
"BF": K_mrad < self.bf_outer,
}
# ── Probe function ────────────────────────────────────────────────────────────
def _focused_probe_k(
N: int,
pixel_size: float,
wavelength: float,
convergence_mrad: float,
defocus_ang: float = 0.0,
cs_mm: float = 0.0,
) -> np.ndarray:
"""
Return the focused probe in k-space (before IFFT to real space).
The aperture function A(k) and CTF phase χ(k) are applied here.
"""
freq = np.fft.fftfreq(N, d=pixel_size)
KX, KY = np.meshgrid(freq, freq, indexing="ij")
K = np.sqrt(KX ** 2 + KY ** 2)
k_c = convergence_mrad * 1e-3 / wavelength
A = (K <= k_c).astype(complex)
k2 = KX ** 2 + KY ** 2
chi = np.pi * wavelength * defocus_ang * k2
if cs_mm:
chi += 0.5 * np.pi * wavelength ** 3 * (cs_mm * 1e7) * k2 ** 2
return A * np.exp(-1j * chi)
def _probe_real(
probe_k: np.ndarray,
shift_x: int = 0,
shift_y: int = 0,
) -> np.ndarray:
"""
Real-space probe at a given scan position via Fourier shift theorem.
shift_x, shift_y are integer pixel offsets from centre.
(Sub-pixel shifts would need a phase ramp; pixels are fine here.)
"""
N = probe_k.shape[0]
freq = np.fft.fftfreq(N)
FX, FY = np.meshgrid(freq, freq, indexing="ij")
# Fourier shift theorem: multiply by exp(-2πi k·r_shift)
phase_shift = np.exp(-2j * np.pi * (FX * shift_x + FY * shift_y))
probe_shifted_k = probe_k * phase_shift
probe_r = np.fft.ifft2(probe_shifted_k)
probe_r /= np.sqrt(np.sum(np.abs(probe_r) ** 2) + 1e-20)
return probe_r
# ── Quantum exit-wave per probe position ──────────────────────────────────────
def _quantum_dwf_exit_wave(
probe_r: np.ndarray,
V_disp: np.ndarray,
sigma: float,
n_q: int,
N: int,
) -> np.ndarray:
"""
Phase grating applied to probe via Qiskit DiagonalGate circuit.
If n_q > MAX_SV_QUBITS, falls back to classical numpy.
"""
if n_q <= MAX_SV_QUBITS:
psi = probe_r.flatten()
psi /= np.linalg.norm(psi) + 1e-20
qc = QuantumCircuit(n_q, name="STEM_Probe")
qc.initialize(psi.tolist(), range(n_q))
grating = np.exp(1j * sigma * V_disp).flatten()
qc.append(DiagonalGate(grating.tolist()), range(n_q))
sv = Statevector.from_instruction(qc)
return sv.data.reshape(N, N)
else:
return probe_r * np.exp(1j * sigma * V_disp)
# ── Core propagation step ─────────────────────────────────────────────────────
def _propagate_to_detector(
psi_exit: np.ndarray,
detector_masks: Dict[str, np.ndarray],
prop_phase: Optional[np.ndarray] = None,
) -> Dict[str, float]:
"""
Forward propagate exit wave and integrate detector signals.
prop_phase : optional Fresnel phase for free-space propagation after specimen.
"""
# Fourier transform to detector (back focal) plane
psi_k = np.fft.fftshift(np.fft.fft2(psi_exit))
if prop_phase is not None:
psi_k = psi_k * np.exp(1j * prop_phase)
I_k = np.abs(psi_k) ** 2
signals = {}
for name, mask in detector_masks.items():
signals[name] = float(np.sum(I_k[mask]))
# iDPC: gradient of phase × BF signal
psi_bfp = np.fft.ifftshift(psi_k)
return signals, I_k, psi_bfp
# ── Main STEM simulation ───────────────────────────────────────────────────────
[docs]
def run_stem(
V: np.ndarray,
pixel_size: float,
voltage: float,
convergence_mrad: float = 25.0,
defocus_ang: float = 0.0,
cs_mm: float = 0.0,
detectors: Optional[STEMDetectors] = None,
scan_step_px: int = 1,
store_4d: bool = False,
) -> Dict:
"""
Run a quantum STEM simulation over the full field of view.
Parameters
----------
V : np.ndarray (N, N)
Projected electrostatic potential [V·Å].
pixel_size : float
Real-space pixel size [Å/pixel].
voltage : float
Accelerating voltage [V].
convergence_mrad : float
Semi-angle of convergence [mrad].
defocus_ang : float
Probe defocus [Å]. Positive → over-focus.
cs_mm : float
Spherical aberration coefficient [mm].
detectors : STEMDetectors or None
Detector configuration. Uses defaults if None.
scan_step_px : int
Scan step in pixels (1 = Nyquist, 2 = half-Nyquist, etc.).
store_4d : bool
If True, store all diffraction patterns → 4D-STEM dataset.
Returns
-------
dict with keys:
'HAADF', 'ADF', 'ABF', 'BF' — STEM images (N_scan, N_scan)
'idpc' — iDPC image
'images' — dict of all images
'KX', 'KY' — k-space axes
'metadata' — parameters dict
'data4d' — 4D array if store_4d=True
"""
if detectors is None:
detectors = STEMDetectors()
N = V.shape[0]
lam = relativistic_wavelength(voltage)
sigma = interaction_constant(voltage, lam)
n_q = 2 * int(np.log2(N))
# Pre-compute probe in k-space (same for all positions)
probe_k = _focused_probe_k(N, pixel_size, lam, convergence_mrad,
defocus_ang, cs_mm)
det_masks = detectors.masks(N, pixel_size, lam)
freq = np.fft.fftshift(np.fft.fftfreq(N, d=pixel_size))
KX, KY = np.meshgrid(freq, freq, indexing="ij")
# Scan positions
scan_coords = list(range(0, N, scan_step_px))
n_scan = len(scan_coords)
# Output images
imgs = {name: np.zeros((n_scan, n_scan)) for name in det_masks}
idpc_x = np.zeros((n_scan, n_scan))
idpc_y = np.zeros((n_scan, n_scan))
data4d = np.zeros((n_scan, n_scan, N, N)) if store_4d else None
fully_quantum = n_q <= MAX_SV_QUBITS
for si, ix in enumerate(scan_coords):
for sj, iy in enumerate(scan_coords):
probe_r = _probe_real(probe_k, shift_x=ix - N // 2,
shift_y=iy - N // 2)
psi_exit = _quantum_dwf_exit_wave(probe_r, V, sigma, n_q, N)
sigs, I_k, psi_bfp = _propagate_to_detector(psi_exit, det_masks)
for name in det_masks:
imgs[name][si, sj] = sigs[name]
# iDPC: first moments of BF disc
bf_mask = det_masks["BF"]
idpc_x[si, sj] = np.sum(KX * I_k * bf_mask) / (np.sum(I_k * bf_mask) + 1e-20)
idpc_y[si, sj] = np.sum(KY * I_k * bf_mask) / (np.sum(I_k * bf_mask) + 1e-20)
if store_4d:
data4d[si, sj] = I_k
# iDPC: divergence of the centre-of-mass field
idpc = np.gradient(idpc_x, axis=0) + np.gradient(idpc_y, axis=1)
idpc -= idpc.mean()
idpc_scale = max(abs(idpc.max()), abs(idpc.min())) + 1e-12
idpc /= idpc_scale
images = dict(imgs)
images["iDPC"] = idpc
result = {
"HAADF": imgs["HAADF"],
"ADF": imgs["ADF"],
"ABF": imgs["ABF"],
"BF": imgs["BF"],
"iDPC": idpc,
"images": images,
"KX": KX, "KY": KY,
"metadata": {
"voltage": voltage,
"wavelength": lam,
"pixel_size": pixel_size,
"convergence_mrad": convergence_mrad,
"defocus_ang": defocus_ang,
"cs_mm": cs_mm,
"engine": "wpoa",
"scan_step_px": scan_step_px,
"N": N,
"n_qubits_total": n_q,
},
"metrics": {
"fully_quantum": fully_quantum,
"approach": (
f"Quantum WPOA STEM ({'statevector' if fully_quantum else 'numpy fallback'})"
),
},
}
if store_4d:
result["data4d"] = data4d
return result