Propagators

This module contains various propagators for light fields

All the functions here are designed for use in an automatic differentiation ptychography model. Each function implements a different propagator.

cdtools.tools.propagators.far_field(wavefront)

Implements a far-field propagator in torch

This accepts a torch tensor, where the last dimension represents the real and imaginary components of the wavefield, and returns the far-field propagated version of it assuming it matches the detector dimensions. It assumes that the propagation is purely far-field, without checking that the geometry is consistent with that assumption.

It also assumes that the real space wavefront is stored in an array [i,j] where i corresponds to the y-axis and j corresponds to the x-axis, with the origin following the CS standard of being in the upper right. The zero frequency component of the propagated wavefield is shifted to the center of the array.

Parameters:

wavefront (torch.Tensor) – The JxNxMx2 stack of complex wavefronts to be propagated

Returns:

propagated – The JxNxMx2 propagated wavefield

Return type:

torch.Tensor

cdtools.tools.propagators.near_field(wavefront, angular_spectrum_propagator)

Propagates a wavefront via the angular spectrum method

This function accepts an 3D torch tensor, where the last dimension represents the real and imaginary components of the wavefield, and returns the near-field propagated version of it. It does this using the supplied angular spectrum propagator, which is a premade phase mask representing the Fourier transform of the kernel for light propagation in the desired geometry.

Parameters:

  • wavefront (torch.Tensor) – The (Leading Dims)xNxM stack of complex wavefronts to be propagated

  • angular_spectrum_propagator (torch.Tensor) – The NxM phase mask to be applied during propagation

Returns:

propagated – The propagated wavefront

Return type:

torch.Tensor

cdtools.tools.propagators.generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, bandlimit=None, **kwargs)

Generates an angular-spectrum based near-field propagator from experimental quantities

This function generates an angular-spectrum based near field propagator that will work on torch Tensors. The function is structured this way - to generate the propagator first - because the generation of the propagation mask is a bit expensive and if this propagator is used in a reconstruction program, then it will be best to calculate this mask once and reuse it.

Formally, this propagator is the the fourier transform of the convolution kernel for light propagation in free space

If the optional bandlimit parameter is set, the propagator will be set to zero beyond an explicit bandlimiting frequency. This is helpful if the propagator will be used in a repeated multiply/propagate framework such as a multislice algorithm, where it helps to prevent aliasing.

Parameters:

  • shape (array) – The shape of the arrays to be propagated

  • spacing (array) – The pixel size in each dimension of the arrays to be propagated

  • wavelength (float) – The wavelength of light to simulate propagation of

  • z (float) – The distance to simulate propagation over

  • bandlimit (float) – Optional, a fraction of the full detector radius beyond which to set the propagator to zero.

Returns:

propagator – A phase mask which accounts for the phase change that each plane wave will undergo.

Return type:

torch.Tensor

cdtools.tools.propagators.inverse_far_field(wavefront)

Implements the inverse of the far-field propagator in torch

This accepts a torch tensor, where the last dimension represents the real and imaginary components of the propagated wavefield, and returns the un-propagated array.

It assumes that the real space wavefront is stored in an array [i,j] where i corresponds to the y-axis and j corresponds to the x-axis, with the origin following the CS standard of being in the upper right. The zero frequency component of the propagated wavefield is assumed to be the center of the array.

Parameters:

wavefront (torch.Tensor) – The JxNxMx2 stack of complex wavefronts propagated to the far-field

Returns:

propagated – The JxNxMx2 exit wavefield

Return type:

torch.Tensor

cdtools.tools.propagators.inverse_near_field(wavefront, angular_spectrum_propagator)

Inverse propagates a wavefront via the angular spectrum method

This function accepts an 3D torch tensor, where the last dimension represents the real and imaginary components of the wavefield, and returns the near-field propagated version of it. It does this using the supplied angular spectrum propagator, which is a premade phase mask.

It propagates the wave using the complex conjugate of the supplied phase mask. This corresponds to propagation backward across the original propagation region - however, the treatment of evanescent waves is such that evanescent waves will decay both during the forward propagation and inverse propagation. This is done for reasons of numerical stability, as the choice to magnify evanescent waves during the inverse propagation process will quickly lead to magnification of any small amount of noise at frequencies larger than k_0, and in most typical situations will even lead to overflow of the floating point range. If evanescent waves need to be treated appropriately for any reason, it is recommended to use the “magnify_evanescent” option in the appropriate helper function used to generate the propagation phase mask. In this case, evanescent waves will be magnified both when used with the forward and inverse near field functions

Parameters:

  • wavefront (torch.Tensor) – The JxNxMx2 stack of complex wavefronts to be propagated

  • angular_spectrum_propagator (torch.Tensor) – The NxM phase mask to be applied during propagation

Returns:

propagated – The inverse propagated wavefront

Return type:

torch.Tensor

cdtools.tools.propagators.generate_high_NA_k_intensity_map(sample_basis, det_basis, det_shape, distance, wavelength, *args, lens=False, **kwargs)

Generates k-space and intensity maps to allow for high-NA far-field propagation of light

At high numerical apertures or for very tilted samples, the simple linear map between location on the detector and location in k-space starts to break down. In addition, at angles above roughly 15 degrees, a correction is needed to account for the decreasing solid angle intersected by each pixel on the detector.

This function generates a map which can be used to apply both corrections via the high_NA_far_field propagator. The k-map which is output is defined as a map between pixel location on the detector and locations in the output of the standard, low-NA far-field propagated wavefield. The output coordinate system is defined to run from -1 to 1 in both directions - this allows for compatibility with pytorch’s grid_sample function. Some detector pixels may be mapped to values outside the rangel [-1,1], depending on the respective sample and detector geometries. This is most likely to be the case if the major correction is due to a tilted sample.

The intensity map is simply an object, the shape of the detector, which encodes intensity corrections between 0 and 1 per pixel.

If the optional “lens” parameter is set to True, the intensity map will be set to a uniform map, and the distortion of Fourier space due to the flat nature of the detector (that is, the portion of the distortion that exists even if the sample is not tilted) will be disabled. This is to account for the fact that a good, infinity-conjugate imaging lens will do it’s best to correct for these abberations in the lens. Of course, the lens will not be perfect, but in such a case it is a better approximation to assume that the lens is perfect than to assume that it is not there at all.

Parameters:

  • sample_basis (array) – The 3x2 sample basis, in real space

  • det_basis (array) – The 3x2 detector basis, in real space

  • det_shape (array) – The length-2 shape of the detector array, (N,M)

  • distance (float) – The sample-to-detector distance

  • wavelength (float) – The wavelength of light being propagated

  • lens (bool) – Whether the diffraction pattern is formed by a lens or not.

Returns:

  • k_map (torch.Tensor) – An NxMx2 tensor mapping detector pixels to locations in the low-NA propagated wavefield

  • intensity_map (torch.Tensor) – An NxM tensor encoding the high-NA intensity correction

cdtools.tools.propagators.high_NA_far_field(wavefront, k_map, intensity_map=None)

Performs a far-field propagation step including a correction for high-NA scenarios

Two major corrections need to be performed when propagating light fields into the far field at high numerical aperture or when the sample is tilted as compared to the detector. The first correction is a deviation from the linear relationship between detector position and spatial frequency in the near field. This is accounted for with the k_map argument, as generated by the generate_high_NA_k_intensity_map function.

The second correction is the change in the solid angle which each pixel subtends at high NA. This is accounted for with an optional intensity map. This is kept optional because some detectors - specifically, those for penetrating radiation - may either not need a correction or need a different correction due to the volumetric nature of the pixels.

If the k-map maps any pixels on the detector to pixels outside of the k-space range of the wavefront, these will be set to zero. This is in keeping with the typical assumption that the sample is band-limited to the Nyquist frequency for the array on which it is sampled.

Parameters:

  • wavefront (torch.Tensor) – The JxNxMx2 stack of complex wavefronts propagated to the far-field

  • k_map (torch.Tensor) – The NxMx2 map accounting for high NA distortion, as generated by generate_high_NA_k_intensity_map

  • intensity_map (torch.Tensor) – The optional NxM tensor accounting for the intensity variation across the detector

Returns:

propagated – The JxNxMx2 exit wavefield

Return type:

torch.Tensor

cdtools.tools.propagators.generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, offset_vector, *args, propagation_vector=None, propagate_along_offset=False, **kwargs)

Generates an angular-spectrum based near-field propagator from experimental quantities

This function generates an angular-spectrum based near field propagator that will work on torch Tensors. The function is structured this way - to generate the propagator first - because the generation of the propagation mask is a bit expensive and if this propagator is used in a reconstruction program, it will be best to calculate this mask once and then reuse it it.

Formally, this propagator is the fourier transform of the convolution kernel for light propagation in free space. It will map a light field at an input plane, with the size and shape defined by the shape and basis inputs, and map it to a plane of the same size and shape offset by the offset vector. It is designed to work on any wavefield defined on an array of parallelograms.

In addition, if the propagation_vector is set, there is an assumed phase ramp applied to the wavefield before propagation, defined such that a feature with uniform phase will propagate along the direction of the propagation vector. This will also remove the phase variation along the propagation direction, because it makes the most physical sense to regard this choice as removing the dominant phase variation in 3D, allowing for the generation of a smoothly varying wavefield over 3D volumes. This decision provides the best numerical stability and allows for the simple setup of light fields copropagating with the coordinate system.

If the propagate_along_offset option is set to True, then the propagation vector will be set equal to the offset vector. This overrides the propagation_vector option

Note that, unlike in the case of the simple angular spectrum propagator, the direction of “forward propagation” is defined by the propagation vector, or (if the propagation vector is not defined), the offset vector. Therefore, in the simple case of a perpendicular offset, there will be no difference between using an offset vector or the negative of the offset vector. This is because, for the light propagation problem to be well posed, the assumption must be made that light only passes through the plane of the known wavefield in one direction. We always assume that light passes through the initial plane travelling in the direction of the final plane.

Practically, if one wants to simulate inverse propagation, there are two possible approaches. First, one can use the inverse_near_field function, which simulates the inverse propagation problem and therefore will naturally simulate propagation in the opposite direction. Second, one can explicitly include a propagation_vector argument, in the direction opposite to the offset vector. Note for both cases that this function will only return propagators which suppress evanescent waves. Thus, propagating forward and then backward by either of these two methods will lead to a supression of evanescant waves. If you need a propagator that will cause evanescent waves to undergo exponential growth, good for you, but this function will not provide it for you.

Parameters:

  • shape (array) – The shape of the arrays to be propagated

  • basis (array) – The (2x3) set of basis vectors describing the array to be propagated

  • wavelength (float) – The wavelength of light to simulate propagation of

  • offset_vector (array) – The displacement to propagate the wavefield along.

  • propagation_vector (array) – The vector along which to include an implied phase ramp to propagate uniform phase features along, if set

  • propagate_along_offset (bool) – Overrides propagation_vector, sets the propagation vector to equal the offset vector if set.

Returns:

propagator – A phase mask which accounts for the phase change that each plane wave will undergo.

Return type:

torch.Tensor