Train a neural phase model on one scan

Start with one measured 4D-STEM counting acquisition. Fit a phase model, compare its output with two direct reconstructions, and save its weights for later scans. The default training patches need at least 128 rows and 256 columns. Use a field large enough for a third disjoint 128 × 128 patch to score it here.

Choose the scan

Edit the two filenames and verify the microscope calibration. Scan pitch is in Å per scan pixel. Use a CUDA kernel with QuantEM.GPU.

from quantem.gpu import DirectPtycho, NeuralPhaseModel, io

scan_path = "data/training_master.h5"
checkpoint = "phase-model.pt"
microscope = dict(voltage_kV=300, semiangle_mrad=30, scan_sampling_A=0.4)  # replace with your measured scan pitch in Å

Fit the optics and train

The direct solver uses the calibrated diffraction data. When the detector sampling is missing, QuantEM estimates it from the bright-field disk. Training splits detected electrons into two halves. A residual that improves prediction of the opposite half earns a better score. A separate patch selects the weights, including the option to keep the starting phase.

scan = io.load(scan_path, device="cuda:0", verbose=False)
ptycho = DirectPtycho(scan, **microscope)
optics = ptycho.aberrations(refine_rotation=True, check_rotation="both", show=False)

model = NeuralPhaseModel.fit(ptycho, optics, recipe="count_split", steps=60, seed=42, checkpoint=checkpoint)
model.report()

Where does the neural phase come from?

The U-Net reads the normalized direct phase and a fixed noise-level channel. It proposes a zero-mean image in radians to add to the phase. A separate learned readout combines local phase features with a map of the counts in a detector annulus. For these settings, the annulus spans 31.5–45 mrad. A bright spot in the annulus can influence the proposal near that position. Its intensity is not itself a phase measurement.

The solver filters both proposals through the current scan’s diffraction equations, then adds a bounded correction based on the remaining mismatch. The last correction tapers to zero at 1.5 Å⁻¹. Training happens once. Applying the model to a new scan uses that scan’s own diffraction and annular counts.

Compare the images

Read the phase-only solution, its neural refinement, and the phase and relative amplitude from the complex fit. The phase panels share contrast limits and report their damping. A scale bar is shown on each image. The second figure pairs neural phase with the measured counts in its guide annulus, on the native scan grid.

series = ptycho.series([scan], fit=optics, method=("phase", "complex"),
                      neural_model=model, refit="none", show=False, verbose=False)
series.show()

reconstruction = series.latest
direct_phase = reconstruction.phase
neural_phase = reconstruction.neural_phase
complex_phase = reconstruction.complex_phase
relative_amplitude = reconstruction.amplitude
annular_counts = reconstruction.annular_counts
reconstruction.show()

Score unused counts

The API picks a third patch outside fitting and selection. Each electron half predicts the other under the same optics. Lower error is better. The complex score uses phase and amplitude together. This measures diffraction prediction, not spatial resolution.

model.score(ptycho, optics)

ptycho.close()
scan.close()