Interactions

This module contains various to simulate stages in the probe-sample interaction

All the tools here are designed to work with automatic differentiation. Each function simulates some aspect of an interaction model that can be used for ptychographic reconstruction.

cdtools.tools.interactions.translations_to_pixel(basis, translations, surface_normal=tensor([0., 0., 1.]))

Takes real space translations and outputs them in pixel space

This works for any 2D ptychography geometry. It takes in A set of translations in (x,y) space and outputs the same translations in internal pixel units perpendicular to the detector.

It uses information on the wavefield basis and, if defined, the sample normal, to perform the conversion.

The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1]. The default sample orientation has a surface normal parallel to this direction

Parameters:

  • basis (torch.Tensor) – The real space basis the wavefields are defined in

  • translations (torch.Tensor) – A Jx3 stack of real-space translations, or a single translation

  • surface_normal (torch.Tensor) – Optional, the sample’s surface normal

Returns:

pixel_translations – A Jx2 stack of translations, or a single translation, in (i,j) pixel-space

Return type:

torch.Tensor

cdtools.tools.interactions.pixel_to_translations(basis, pixel_translations, surface_normal=tensor([0., 0., 1.]))

Takes pixel-space translations and outputs them in real space

This works for any 2D ptychography geometry. It takes in A set of internal pixel unit translations in (i,j) space and outputs the same translations real (x,y) space

It uses information on the wavefield basis and, if defined, the sample normal, to perform the conversion.

The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1]. The default sample orientation has a surface normal parallel to this direction. Because of this, the z direction translation is always set to zero in the conversion

Parameters:

  • basis (torch.Tensor) – The real space basis the wavefields are defined in

  • translations (torch.Tensor) – A Jx2 stack of pixel-space translations, or a single translation

  • surface_normal (torch.Tensor, default: torch.Tensor([0,0,1])) – The sample’s surface normal

Returns:

real_translations – A Jx3 stack of real-space translations, or a single translation

Return type:

torch.Tensor

cdtools.tools.interactions.project_translations_to_sample(sample_basis, translations)

Takes real space translations and outputs them in pixels in a sample basis

This projection function is designed for the Bragg2DPtycho class. More broadly, it works to take a set of translations in the lab frame and convert each one into two values. First, an (i,j) value in pixels describing the location of the probe’s intersection with the sample plane, assuming the basis found in sample_basis is used. Second, an amount that the probe must be propagated along the z-axis to reach the sample plane at the given location. This includes both the effect of the tilted sample plane and any explicitly defined motion along the z axis of the probe-forming optic as included in the input translations.

Note that because of the sign convention (that this function returns th relative amount the probe needs to be propagated to reach any given location), a positive motion along the z-axis of the probe forming optics will lead to a negative propagation distance.

The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1].

Parameters:

  • sample_basis (torch.Tensor) – The real space basis the wavefields are defined in

  • translations (torch.Tensor) – A Jx3 stack of real-space translations, or a single translation

Returns:

  • pixel_translations (torch.Tensor) – A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation

  • propagations (torch.Tensor) – A length-J vector of propagation distances, in meters

cdtools.tools.interactions.ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False)

Returns a stack of exit waves without accounting for subpixel shifts

This function returns a collection of exit waves, with the first dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by shifting the probe by the rounded value of the translation

If multiple_modes is set to False, any additional dimensions in the ptycho_2D_round function will be assumed to correspond to the translation index. If multiple_modes is set to true, the (-4th) dimension of the probe will always be assumed to be defining a set of (P) incoherently mixing modes to be broadcast all translation indices. If any additional dimensions closer to the start exist, they will be assumed to be translation indices

Parameters:

  • probe (torch.Tensor) – A (P)xMxL probe function to illuminate the object

  • object (torch.Tensor) – The object function to be probed

  • translations (torch.Tensor) – The (N)x2 array of (i,j) translations to simulate

  • multuple_modes (bool) – Default False, whether to assume the probe contains multiple modes

Returns:

exit_waves – An (N)x(P)xMxL tensor of the calculated exit waves

Return type:

torch.Tensor

cdtools.tools.interactions.ptycho_2D_linear(probe, obj, translations, shift_probe=True)

Returns a stack of exit waves accounting for subpixel shifts

This function returns a collection of exit waves, with the first dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by shifting the probe with each translation in turn, using linear interpolation to combine the results

If shift_probe is True, it applies the subpixel shift to the probe, otherwise the subpixel shift is applied to the object

Parameters:

  • probe (torch.Tensor) – An MxL probe function for the exit waves

  • object (torch.Tensor) – The object function to be probed

  • translations (torch.Tensor) – The Nx2 array of translations to simulate

  • shift_probe (bool) – Default True, Whether to subpixel shift the probe or object

Returns:

exit_waves – An NxMxL tensor of the calculated exit waves

Return type:

torch.Tensor

cdtools.tools.interactions.ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, probe_support=None)

Returns a stack of exit waves accounting for subpixel shifts

This function returns a collection of exit waves, with the first dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by shifting the probe with each translation in turn, using sinc interpolation (done via multiplication with a complex exponential in Fourier space)

If shift_probe is True, it applies the subpixel shift to the probe, otherwise the subpixel shift is applied to the object [not yet implemented]

If multiple_modes is set to False, any additional dimensions in the ptycho_2D_round function will be assumed to correspond to the translation index. If multiple_modes is set to true, the (-4th) dimension of the probe will always be assumed to be defining a set of (P) incoherently mixing modes to be broadcast all translation indices. If any additional dimensions closer to the start exist, they will be assumed to be translation indices

Parameters:

  • probe (torch.Tensor) – An (P)xMxL probe function for the exit waves

  • object (torch.Tensor) – The object function to be probed

  • translations (torch.Tensor) – The (N)x2 array of translations to simulate

  • shift_probe (bool) – Default True, Whether to subpixel shift the probe or object

  • multuple_modes (bool) – Default False, whether to assume the probe contains multiple modes

Returns:

exit_waves – An (N)x(P)xMxL tensor of the calculated exit waves

Return type:

torch.Tensor

cdtools.tools.interactions.RPI_interaction(probe, obj)

Returns an exit wave from a high-res probe and a low-res obj

In this interaction, the probe and object arrays are assumed to cover the same physical region of space, but with the probe array sampling that region of space more finely. Thus, to do the interaction, the object is first upsampled by padding it in Fourier space (equivalent to a sinc interpolation) before being multiplied with the probe. This function is called RPI_interaction because this interaction is central to the RPI method and is not commonly used elsewhere.

This also works with object functions that have an extra first dimension for an incoherently mixing model.

Parameters:

  • probe (torch.Tensor) – An MxL probe function for simulating the exit waves

  • obj (torch.Tensor) – An M’xL’ or NxM’xL’ object function for simulating the exit waves

Returns:

exit_wave – An MxL tensor of the calculated exit waves

Return type:

torch.Tensor