Data structures

Pty-Chi is built on PyTorch, so on the most basic level, all data structures are PyTorch tensors. When a hardware accelerator device (like a GPU) is used, these tensors might reside on the device. To plot them or save them to disk, transfer them to the CPU using tensor.cpu(). You may also convert them to NumPy arrays using tensor.cpu().numpy().

For user friendliness, Pty-Chi’s high-level interface accepts both NumPy arrays and PyTorch tensors for initial guesses and data for the object, probe, probe positions, and initial OPR mode weights. If PyTorch tensors are used, they do not need to be transferred to accelerator devices as Pty-Chi does this internally.

Tensor shapes

The table below lists the expected shapes of the tensors for different reconstruction parameters.

Parameter

Shape

Remarks

Object

(n_slices, height, width)

Probe

(n_opr_modes, n_incoherent_modes, height, width)

Probe positions

(n_positions, 2)

  1. Probe positions should follow row-major order, i.e., y-coordinates come first.

  2. The aforementioned shape is of the internally stored position tensor. The ProbePositionOptions object takes x- and y-positions as separate arguments.

Diffraction data

(n_positions, height, width)

OPR mode weights

(n_positions, n_opr_modes)

Fetching data

The high-level Python interface of Pty-Chi manages reconstruction jobs using the PtychographyTask class. When the reconstruction job finishes or is paused, you may get the data using get_data_to_cpu(). For example:

# Get the object
object = task.get_data_to_cpu('object', as_numpy=True)

Alternatively, you can also get reconstruction parameters from the task object directly without data structure conversion. Reconstruction parameters, such as the object and the probe, are stored in task.parameter_group. By calling task.parameter_group.object, you get a ReconstructionParameter object for the object function. You can then get the data from the parameter object using .data.