cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-10 05:22:41 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
aea0ff9bf8097d0b2cf7f4b59d69e3cfecb6611b
cdtools/CDTools/tools/propagators.py
T

646 lines
26 KiB
Python
Raw Blame History

"""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.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
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']
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
"""
return fftshift(t.fft(ifftshift(wavefront), 2, normalized=True))
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
"""
return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
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)
#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 map 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):
real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
result = t.stack((real_output[0,:,:,:],imag_output[0,:,:,:]),dim=3)
if intensity_map is not None:
result = result * intensity_map[None,:,:,None]
return result
original_dim = wavefront.dim()
if original_dim == 3:
result = process_wavefield_stack(low_NA_wavefield[None,:,:,:])
return result[0,:,:,:]
if original_dim == 4:
result = process_wavefield_stack(low_NA_wavefield)
return result
if original_dim == 5:
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, **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 complex conjugate of the fourier
transform of the convolution kernel for light propagation in free
space
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
Returns
-------
propagator : torch.Tensor
A phase mask which accounts for the phase change that each plane wave will undergo.
"""
ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0])
kj = 2 * np.pi * fftpack.fftfreq(shape[1],spacing[1])
Kj, Ki = np.meshgrid(kj,ki)
# Define this as complex so the square root properly gives
# k>k0 components imaginary frequencies
k0 = np.complex128((2*np.pi/wavelength))
# Properly accuount for evanescent waves
if z >=0:
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z)
else:
propagator = np.exp(1j*np.conj(np.sqrt(k0**2 - Ki**2 - Kj**2)) * z)
if remove_z_phase:
propagator *= np.exp(-1j * k0 * z)
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
propagator = complex_to_torch(np.conj(propagator))
return propagator.to(*args, **kwargs)
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, then it will be best
to calculate this mask once and close over it.
Formally, this propagator is the complex conjugate of the fourier
transform of the convolution kernel for light propagation in free
space. It will map a ligh 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 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. Mathematically, this
corresponds to a choice of uniform phase objects either accumulating
positive or negative phase. In the simple propagation case, there is
no ambiguity introduced by always choosing the light field to propagate
along the positive z direction. In the general case, there is no equivalent
obvious choice - thus, the light is always assumed to pass through the
initial plane travelling in the direction of the final plane.
Practically, if one wants to simulate inverse propagation, there are then
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, which overrides
the offset vector in defining the direction in which light passes through
the input plane. However, in this case, the resulting light field will have
the overall phase accumulation due to propagation along the propagation
vector removed, which may not be the intended behavior. However, this is
not recommended, as inverse propagation will tend to magnify evanescent
waves - it is therefore preferable (unless there is a specific need to
account for evanescent waves properly) to use the inverse near field
propagator
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.
"""
# These check for any pytorch inputs and convert them
try:
basis = basis.detach().cpu().numpy()
except:
pass
shape = np.array(tuple(shape))
try:
offset_vector = offset_vector.detach().cpu().numpy()
except:
pass
if propagation_vector is not None:
try:
propagation_vector = propagation_vector.detach().cpu().numpy()
except:
pass
try:
wavelength = wavelength.detach().cpu().numpy()[0]
except:
pass
# First we calculate a dual basis for the real space grid
inv_basis = np.linalg.pinv(basis).transpose()
# Then we calculate the frequencies in (i,j) space
ki = 2 * np.pi * fftpack.fftfreq(shape[0])
kj = 2 * np.pi * fftpack.fftfreq(shape[1])
K_ij = np.stack(np.meshgrid(ki,kj, indexing='ij'))
# Now we convert these to frequencies in reciprocal space
# These frequencies span the 2D plane of the input wavefield.
K_xyz = np.tensordot(inv_basis, K_ij, axes=1)
# Now we need to apply two corrections to the standard AS method.
# First, we calculate a phase mask which corresponds to the
# shift of the final plane away from the perpendicular direction
# from the input plane. We don't need to extract the perpendicular
# component of the shift because the K_xyz vectors are naturally in the
# input plane.
# This may have a sign error - must be checked
phase_mask = np.exp(-1j * np.tensordot(offset_vector,K_xyz,axes=1))
# Next, we apply a shift to the k-space vectors which sets up
# propagation such that a uniform phase object will propagate along the
# propagation axis. This is not modeling a physical effect, but simply is
# the clearest way to do a rigorous simulation while preventing
# aliasing-related challenges. If used (as is by default), be aware
# and prepare the input wavefields appropriately.
if propagate_along_offset:
propagation_vector = offset_vector
perpendicular_dir = np.cross(basis[:,1],basis[:,0])
perpendicular_dir /= np.linalg.norm(perpendicular_dir)
offset_perpendicular = np.dot(perpendicular_dir, offset_vector)
k0 = 2*np.pi/wavelength
sign_correction = 1
# Only implement the shift if the flag is set to True
if propagation_vector is not None:
prop_perpendicular = np.dot(perpendicular_dir, propagation_vector)
prop_parallel = propagation_vector - perpendicular_dir \
* prop_perpendicular
if np.linalg.norm(propagation_vector) == 0:
# for numerical stability if this is exactly zero we need
# a special case
k_offset = np.array([0,0,0])
else:
k_offset = prop_parallel * k0 / np.linalg.norm(propagation_vector)
K_xyz = K_xyz - k_offset[:,None,None]
# There apparently is a sign correction that I need to apply
#sign_correction = np.sign(np.dot(perpendicular_dir,propagation_vector))
sign_correction = np.sign(np.dot(offset_vector,propagation_vector))
# we also need to remove the z-dependence on the phase
# This time, though, the z-dependence actually has to do with
# the out of plane component of k at the central offset. Normally
# this is 0, so the z-component is just k0, but not in this case
# We only need one case here, unlike with the propagator, because
# k_offset will always be less than k0
phase_mask *= np.exp(-1j * np.sqrt(k0**2 - np.linalg.norm(k_offset)**2)
* sign_correction
* np.abs(offset_perpendicular))
# Redefine this as complex so the square root properly gives
# k>k0 components imaginary frequencies
k0 = np.complex128(k0)
# Finally, generate the propagator!
# Must have cases to ensure that evanescent waves decay instead of grow
if sign_correction > 0:
propagator = np.exp(1j*np.sqrt(k0**2 - np.linalg.norm(K_xyz,axis=0)**2)
* sign_correction * np.abs(offset_perpendicular))
else:
propagator = np.exp(-1j * np.conj(np.sqrt(k0**2 -
np.linalg.norm(K_xyz,axis=0)**2))
* np.abs(offset_perpendicular))
propagator *= phase_mask
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
propagator = complex_to_torch(np.conj(propagator))
return propagator.to(**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.
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 propagated wavefront
"""
return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2)
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.ifft(cmult(t.fft(wavefront,2), cconj(angular_spectrum_propagator)), 2)
# 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
Reference in New Issue View Git Blame Copy Permalink