"""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. """ import torch as t from torch.nn.functional import grid_sample from scipy import fftpack import numpy as np from matplotlib import pyplot as plt __all__ = ['far_field', 'near_field', 'generate_angular_spectrum_propagator', 'inverse_far_field', 'inverse_near_field', 'generate_high_NA_k_intensity_map', 'high_NA_far_field', 'generate_generalized_angular_spectrum_propagator'] def 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 : torch.Tensor The JxNxMx2 propagated wavefield """ shifted = t.fft.ifftshift(wavefront, dim=(-1,-2)) propagated = t.fft.fft2(shifted, norm='ortho') return t.fft.fftshift(propagated, dim=(-1,-2)) def 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 : torch.Tensor The JxNxMx2 exit wavefield """ shifted = t.fft.ifftshift(wavefront, dim=(-1,-2)) propagated = t.fft.ifft2(shifted, norm='ortho') return t.fft.fftshift(propagated, dim=(-1,-2)) def 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 """ # First we convert things to numpy try: sample_basis = sample_basis.detach().cpu().numpy() except: pass try: det_basis = det_basis.detach().cpu().numpy() except: pass det_shape = np.array(tuple(det_shape)) try: distance = distance.detach().cpu().numpy()[0] except: pass try: wavelength = wavelength.detach().cpu().numpy()[0] except: pass # The next order of business is to calculate the k values associated with # each pixel. i_arr = np.arange(det_shape[0]) i_arr = i_arr - np.mean(i_arr) j_arr = np.arange(det_shape[1]) j_arr = j_arr - np.mean(j_arr) Is, Js = np.meshgrid(i_arr,j_arr,indexing='ij') samp_det_vec = np.cross(det_basis[:,0],det_basis[:,1]) samp_det_vec *= distance / np.linalg.norm(samp_det_vec) # This could potentially correct for a mistake in the implied # propagation direction (e.g. choosing e^ikx instead of e^-ikx) # This appears to be correct, based on empirical evidence from # a grazing incidence reflection experiment at 10 degrees # on the optical table samp_det_vec *= -1 if lens == False: # This correctly reproduces the sample-to-each-pixel vectors # in the case where the diffraction pattern is actually formed # by Fraunhoffer diffraction Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1) \ + samp_det_vec[:,None,None] else: # This forms a distorted set of vectors designed to produce the # correct Fourier space map in the case where an imaging lens is # used in the 2f geometry. One should not read too much meaning # into these vectors, they are simply set up to produce the # correct final K-map Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1)# Rs += (samp_det_vec / np.linalg.norm(samp_det_vec))[:,None,None] * \ np.sqrt(np.sum((samp_det_vec)**2)-np.sum(Rs**2,axis=0))[None,:,:] k0 = 2*np.pi/wavelength Ks = k0 * Rs / np.linalg.norm(Rs, axis=0) # My attempt at seeing what happens if I flip the Ks #Ks *= -1 # This is the cosine of the angle with the detector normal intensity_map = np.tensordot(samp_det_vec/(k0*distance),Ks,axes=1) if lens: # Set the intensity map to be uniform if a lens is being used intensity_map = np.ones_like(intensity_map) intensity_map = t.Tensor(intensity_map).to(*args, **kwargs) # This accounts for the implied phase ramp along the exit wave direction # In other words, it prevents the diffraction pattern from sliding off the # detector when the sample is tilted but represented by an object with # uniform phase. Ks -= k0 * samp_det_vec[:,None,None] / distance # A potential alternative when Ks are flipped #Ks += k0 * samp_det_vec[:,None,None] / distance # Now we move on to finding the conversion into k-space # for the sample grid. It turns out we can do this by multiplying # them with the real space basis (dual of the reciprocal space # basis is the real space basis). In fact, because we want to return # values scaled to the overall size of the k-space window, we don't # even need the shape of the sample array k_map = np.tensordot(2*sample_basis.transpose()[::-1,:] / (2*np.pi),Ks,axes=1) k_map = t.Tensor(np.moveaxis(k_map,0,2)).to(*args, **kwargs) # Potentially we need a correction to account for the discrete nature # of the FFT # And finally, we need to convert the results to pytorch return k_map, intensity_map def 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 : torch.Tensor The JxNxMx2 exit wavefield """ low_NA_wavefield = far_field(wavefront) # I'm going to need to separately interpolate the real and complex parts # This can be done k_map = k_map[None,:,:,:] # Will only work for a 4D wavefile stack. #plt.figure() #plt.pcolormesh(k_map[0,:,:,0].cpu().numpy(),k_map[0,:,:,1].cpu().numpy(), # np.ones_like(k_map[0,:-1,:-1,0].cpu().numpy())) #plt.show() def process_wavefield_stack(low_NA_wavefield): # grid_sample doesn't work on complex-valued wavefields real_output = grid_sample(low_NA_wavefield[None,:,:,:].real,k_map,mode='bilinear',padding_mode='zeros', align_corners=False) imag_output = grid_sample(low_NA_wavefield[None,:,:,:].imag,k_map,mode='bilinear',padding_mode='zeros', align_corners=False) result = real_output[0,:,:,:] + 1j * imag_output[0,:,:,:] if intensity_map is not None: result = result * intensity_map[None,:,:] return result original_dim = wavefront.dim() if original_dim == 2: result = process_wavefield_stack(low_NA_wavefield[None,:,:]) return result[0,:,:] if original_dim == 3: result = process_wavefield_stack(low_NA_wavefield) return result if original_dim == 4: result = [] for i in range(low_NA_wavefield.size()[0]): result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:])) return t.stack(result) else: raise IndexError('Wavefield had incorrect number of dimensions') def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, 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 remove_z_phase : bool Default False, whether to remove the dominant z-direction phase dependence bandlimit : float Optional, a fraction of the full detector radius beyond which to set the propagator to zero. Returns ------- propagator : torch.Tensor A phase mask which accounts for the phase change that each plane wave will undergo. """ # Internally, the generalized propagation function is used, so we start # by creating an appropriate basis # This creates a real-valued tensor which matches the kind of complex # number dtype requested in **kwargs if 'dtype' in kwargs: basis = t.real(t.zeros([3,2], **kwargs)) else: basis = t.zeros([3,2], dtype=t.float32) spacing = t.as_tensor(spacing, dtype=basis.dtype) basis[1,0] = -spacing[0] basis[0,1] = -spacing[1] # And similarly, the offset is just z along the z direction #offset = t.tensor([0,0,z], dtype=basis.dtype) offset = t.zeros(3, dtype=basis.dtype) offset[2] = z # And we call the generalized function! propagator = generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, offset, propagate_along_offset=remove_z_phase, **kwargs) if z < 0: propagator = t.conj(propagator) # Bandlimiting is not implemented in the generalized function, because it # has a less clear meaning in that setting, so we apply it here instead if bandlimit is not None: # No need to multiply by 2pi ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0]) kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1]) Ki, Kj = t.meshgrid(ki,kj) min_radius = min(t.max(ki),t.max(kj)) Rs = t.sqrt((Ki/t.max(ki))**2 + (Kj/t.max(kj))**2) propagator = propagator * (Rs < bandlimit) return propagator def 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 : torch.Tensor A phase mask which accounts for the phase change that each plane wave will undergo. """ # make sure everything is in pytorch, and set the propagation vector # appropriately if propagate_along_offset is chosen basis = t.as_tensor(basis) offset_vector = t.as_tensor(offset_vector, dtype=basis.dtype) if propagate_along_offset: propagation_vector = offset_vector if propagation_vector is not None: propagation_vector = t.as_tensor(propagation_vector, dtype=basis.dtype) # # In this section, we calculate the wavevectors associated with each # pixel in Fourier space. This is the meat of the function # # First we calculate a dual basis for the real space grid inv_basis = t.linalg.pinv(basis).transpose(0,1) # Then we calculate the frequencies in (i,j) space ki = 2 * np.pi * t.fft.fftfreq(shape[0], dtype=inv_basis.dtype) kj = 2 * np.pi * t.fft.fftfreq(shape[1], dtype=inv_basis.dtype) K_ij = t.stack(t.meshgrid(ki,kj)) # Now we convert these to frequencies in reciprocal space # These frequencies span the 2D plane of the input wavefield, # hence K_ip for "in-plane" K_ip = t.tensordot(inv_basis, K_ij, dims=1) # Now, we need to generate the out-of-plane direction, so we can # expand these Ks to the full Ks in 3D reciprocal space. # This is broken down into 2 steps to avoid floating point underflow # which was a real problem that showed up for electron ptycho b1_dir = basis[:,0] / t.linalg.norm(basis[:,0]) b2_dir = basis[:,1] / t.linalg.norm(basis[:,1]) perpendicular_dir = t.cross(b1_dir,b2_dir) # Note that we cannot use in-place operations if we want to be able to # use automatic differentiation successfully perpendicular_dir = perpendicular_dir / t.linalg.norm(perpendicular_dir) # We set the sign of the propagation direction appropriately if propagation_vector is not None: perpendicular_dir = perpendicular_dir \ * t.sign(t.dot(perpendicular_dir,propagation_vector)) else: pass perpendicular_dir = perpendicular_dir * \ t.sign(t.dot(perpendicular_dir,offset_vector)) # Then, if we have a propagation vector, we shift the in-plane # components of all the pixels to be centered around the in-plane # component of the propagation vector. if propagation_vector is not None: prop_dir = (propagation_vector / t.linalg.norm(propagation_vector)) K_0 = 2*np.pi / wavelength * prop_dir K_0_ip = K_0 - t.dot(perpendicular_dir,K_0) * perpendicular_dir K_ip += K_0_ip[:,None,None] else: K_0 = t.zeros_like(offset_vector) # Now, we have accurate in-plane values for K, so we can calculate the # out-of-plane part. We start by calculating it's squared magnitude K_oop_squared = (2*np.pi/wavelength)**2 - t.linalg.norm(K_ip,dim=0)**2 # Then, we take the square root and assign it the appropriate direction, # adding to get the full 3D wavevectors. Note that we convert to complex # before the square root to appropriately map negative numbers to # complex frequencies K = K_ip + perpendicular_dir[:,None,None] \ * t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared))) # # In this section, we take the inner product of the calcualted # wavevectors with the offset vector, to get the phase shift # experienced by each plane wave. # # We need to convert to complex because K is complex offset_vector = t.complex(offset_vector,t.zeros_like(offset_vector)) # We then subtract off K_0, essentially setting the phase offset # experienced by K_0 to 0. K_0 will already have been set to 0 if # there was no propagation vector set. K_m_K_0 = K - K_0[:,None,None] # We actually calculate the phase mask phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1) # And we take the conjugate if needed, which makes sure that the # imaginary part is always positive (supressing evanescent waves) if propagation_vector is not None \ and t.sign(t.dot(propagation_vector,offset_vector.real)) == -1: phase_mask = t.conj(phase_mask) # We return the final mask, in the form asked for! return t.exp(1j*phase_mask).to(*args,**kwargs) def 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 : torch.Tensor The propagated wavefront """ return t.fft.ifft2(angular_spectrum_propagator * t.fft.fft2(wavefront)) def 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 : torch.Tensor The inverse propagated wavefront """ return t.fft.ifft2(t.fft.fft2(wavefront) * t.conj(angular_spectrum_propagator)) # I think it would be worthwhile to implement an FFT-DI based strategy as # well, especially for probe initialization where the propagation distance # can be large relative to what the angular spectrum method can reliably handle