ptyrad.core.forward

ptyrad.core.forward#

Physical forwad model that generates diffraction patterns from mixed-state probe/object in a fully vectorized way

Functions

multislice_forward(obja_patches, ...[, ...])

Computes the multislice electron diffraction pattern with multiple incoherent probe and object modes using a vectorized forward model.
ptyrad.core.forward.multislice_forward(obja_patches, objp_patches, probe, H, omode_occu=None, eps=1e-10)[source]#

Computes the multislice electron diffraction pattern with multiple incoherent probe and object modes using a vectorized forward model.

Parameters:

  • obja_patches (torch.Tensor) – Tensor of shape (N, omode, Nz, Ny, Nx), representing object amplitude patches with float32. N is the number of samples in a batch, omode is the number of object modes, Nz, Ny, Nx are the dimensions of the object patches.

  • objp_patches (torch.Tensor) – Tensor of shape (N, omode, Nz, Ny, Nx), representing object phase patches with float32. N is the number of samples in a batch, omode is the number of object modes, Nz, Ny, Nx are the dimensions of the object patches.

  • omode_occu (torch.Tensor) – Tensor of shape (omode,) with float32 values, representing the occupancy/expectation for each object mode. The sum of all elements should be 1.

  • probe (torch.Tensor) – Tensor of shape (N, pmode, Ny, Nx) with complex64 values, representing the probe(s). N is the number of samples in the batch, pmode is the number of probe modes. By default, N is 1, assuming the same probe for all samples.

  • H (torch.Tensor) – Tensor of shape (N, Ky, Kx) with complex64 values, representing the Fresnel propagator that propagates the wave by a slice thickness.

  • eps (float, optional) – A small value added for numerical stability. Defaults to 1e-10.

Returns:

Tensor of shape (N, Ky, Kx) with float32 positive values, representing the forward diffraction pattern for each sample in the batch.

Return type:

torch.Tensor

Contents