mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
242 lines
8.6 KiB
Python
242 lines
8.6 KiB
Python
"""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 scipy import fftpack
|
|
import numpy as np
|
|
|
|
__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_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, **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
|
|
|
|
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
|
|
|
|
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))
|
|
|
|
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * 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, propagation_vector, *args, **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
|
|
|
|
This function is written to work on any wavefield defined on any
|
|
array of parallelograms. In addition, 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
|
|
defined propagation vector. This helps simplify
|
|
|
|
|
|
Parameters
|
|
----------
|
|
shape : array
|
|
The shape of the arrays to be propagated
|
|
spacing : array
|
|
The (2x3) set of basis vectors describing the array to be propagated
|
|
wavelength : float
|
|
The wavelength of light to simulate propagation of
|
|
propagation_vector : array
|
|
The displacement to propagate the wavefield along.
|
|
tilt : float
|
|
The tilt, in radians, of the plane that the wavefield is defined on
|
|
|
|
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))
|
|
|
|
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * 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(**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 conjugate of the supplied phase mask,
|
|
which corresponds to the inverse propagation problem.
|
|
|
|
|
|
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
|
|
|
|
|