Commit ad301752 authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

Major rewrite of vacuum propagation code, including padding.

parent 9763a854
# This program is public domain
"""
Diffraction / Wavefield propagation through vacuum and homogeneous media.
"""
import numpy as np
from scipy.fft import fftn, ifftn
from scipy.signal import fftconvolve
from . import hankel
from .misc import fftfreqn, gridn
from . import hankel
from .misc import fftfreqn, gridn, crop, squaresum
def fresnelKernelCS(N, fresnelNumber):
hFreq = hankel.hankelFreq(N)
kern = np.exp(-1j * np.pi / fresnelNumber * hFreq ** 2)
class FresnelPropagatorCS:
def __init__(self, nsamples, fresnel_numbers, npad=2):
return kern
self._nsamples = nsamples
self._npad = npad
if np.any(self._npad < 1):
raise ValueError("Padding factor `npad` cannot be less than 1.")
class FresnelPropagatorCS:
def __init__(self, N, fresnelNumber):
fresnel_numbers = np.asarray(fresnel_numbers)
if np.ndim(fresnel_numbers) == 0:
self._ndist = 1
elif np.ndim(fresnel_numbers) == 1:
self._ndist = fresnel_numbers.shape[0]
else:
raise ValueError(
"Parameter 'fresnel_numbers' must either be scalar or one-dimensional."
)
self._N = N
self.fresnel_numbers = fresnel_numbers * np.ones((self._ndist,))
self._kern = fresnelKernelCS(N, fresnelNumber)
self._dht = hankel.DiscreteHankelTransform(N, 0)
# calculate amount of padding
self._pad_width = (self._nsamples * (self._npad - 1) / 2).astype(int)
def __call__(self, u):
self._kern = self._init_kernel()
self._dht = hankel.DiscreteHankelTransform(nsamples, 0)
hf_u = self._dht(u)
hf_u *= self._kern
uprop = self._dht(hf_u)
def _init_kernel(self):
samples_pad = self._nsamples + 2 * self._pad_width
return uprop
freq = hankel.hankel_freq(samples_pad)
kern = np.exp(-1j * np.pi / self.fresnel_numbers * np.square(freq))
return kern
def fresnelTFKernel(shape, fresnelNumbers):
def __call__(self, u):
u = np.asarray(u)
ndim = len(shape)
if np.ndim(u) != 1 or u.shape[0] != self._nsamples:
raise ValueError("Invalid shape.")
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(ndim)
# pad with zeros
upad = np.pad(u, self._pad_width)
f = fftfreqn(shape)
kernel = np.ones(shape, dtype=np.complex128)
for dim in range(ndim):
kernel *= np.exp(-1j * np.pi / fresnelNumbers[dim] * f[dim] ** 2)
hf_u = self._dht(upad)
hf_uprop = self._kern * hf_u
uprop = np.zeros_like(u, shape=(self._ndist, *u.shape))
return kernel
for dist in self._ndist:
tmp = self._dht(hf_uprop[dist, :])
uprop[dist, :] = crop(tmp, (self._pad_width, self._pad_width))
return uprop
class FresnelTFPropagator:
def __init__(self, shape, fresnelNumbers):
self._shape = np.array(shape)
class FTConvolutionPropagator:
"""
Blub
"""
def __init__(self, shape, npad):
self._shape = shape
self._ndim = len(self._shape)
self._npad = np.asarray(npad) * np.ones((self._ndim,))
if np.any(self._npad < 1):
raise ValueError("Padding factor `npad` cannot be less than 1.")
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(self._ndim)
# calculate amount of padding
self._pad_width = (np.asarray(self._shape) * (self._npad - 1) / 2).astype(int)
self._kernel = fresnelTFKernel(shape, fresnelNumbers)
self.kernel = 1
def __call__(self, u, workers=-1):
u = np.asarray(u)
if u.shape != self._shape:
raise ValueError("Invalid shape.")
# pad with zeros
upad = np.pad(u, self._pad_width)
ft_u = fftn(u, workers=workers)
ft_u *= self._kernel
uprop = ifftn(ft_u, workers=workers)
ft_u = fftn(upad, axes=(-2, -1), workers=workers)
ft_uprop = ft_u * self.kernel
uproppad = ifftn(ft_uprop, axes=(-2, -1), workers=workers)
# crop central part
pad_width_full = [
0,
] + list(self._pad_width)
uprop = crop(uproppad, zip(pad_width_full, pad_width_full))
return uprop
def fresnelIRKernel(shape, fresnelNumbers):
class FresnelTFPropagator(FTConvolutionPropagator):
"""
Blub
"""
ndim = len(shape)
def __init__(self, shape, fresnel_numbers, npad=2):
super(FresnelTFPropagator, self).__init__(shape, npad)
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(ndim)
fresnel_numbers = np.asarray(fresnel_numbers)
if np.ndim(fresnel_numbers) == 0:
self._ndist = 1
else:
self._ndist = fresnel_numbers.shape[0]
kernel = np.ones(2 * np.array(shape, dtype=int) - 1, dtype=np.complex128)
self.fresnel_numbers = fresnel_numbers * np.ones((self._ndist, self._ndim))
# phase factor
kernel *= np.prod(
np.sqrt(np.abs(fresnelNumbers)) * np.exp((-1j * np.pi / 4) * np.sign(fresnelNumbers))
)
self.kernel = self._init_kernel()
x = gridn(shape)
def _init_kernel(self):
shape_pad = np.asarray(self._shape) + 2 * self._pad_width
f = fftfreqn(shape_pad)
for dim in range(ndim):
kernel *= np.exp(1j * np.pi * fresnelNumbers[dim] * x[dim] ** 2)
kernel = np.ones((self._ndist, *shape_pad), dtype=np.complex128)
for dist in range(self._ndist):
for dim in range(self._ndim):
kernel[dist, :] *= np.exp(
-1j * np.pi / self.fresnel_numbers[dist, dim] * np.square(f[dim])
)
return kernel
return kernel
class FresnelIRPropagator:
def __init__(self, shape, fresnelNumbers):
class FresnelIRPropagator(FTConvolutionPropagator):
"""
Blub
"""
self._shape = np.array(shape)
self._ndim = len(self._shape)
def __init__(self, shape, fresnel_numbers, npad=2):
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(self._ndim)
super(FresnelIRPropagator, self).__init__(shape, npad)
self._kernel = fresnelIRKernel(shape, fresnelNumbers)
fresnel_numbers = np.asarray(fresnel_numbers)
if np.ndim(fresnel_numbers) == 0:
self._ndist = 1
else:
self._ndist = fresnel_numbers.shape[0]
self.fresnel_numbers = fresnel_numbers * np.ones((self._ndist, self._ndim))
def __call__(self, u):
self.kernel = self._init_kernel()
uprop = fftconvolve(u, self._kernel, mode="valid")
def _init_kernel(self):
shape_pad = np.asarray(self._shape) + 2 * self._pad_width
return uprop
# 1. assemble convolution kernel
conv_kernel = np.ones((self._ndist, *shape_pad), dtype=np.complex128)
x = gridn(shape_pad)
for dist in range(self._ndist):
# phase factor
conv_kernel *= np.prod(
np.sqrt(np.abs(self.fresnel_numbers[dist, :]))
* np.exp((-1j * np.pi / 4) * np.sign(self.fresnel_numbers[dist, :]))
)
# chirp function
for dim in range(self._ndim):
conv_kernel *= np.exp(
1j * np.pi * self.fresnel_numbers[dist, dim] * np.square(x[dim])
)
# 2. compute TF from convolution kernel
kernel = fftn(conv_kernel, axes=(-2, -1))
return kernel
class ASMPropagator(FTConvolutionPropagator):
"""
Blub
"""
_lambda0 = 1.0 # all lengths are in units of the wavelength
_k0 = 2 * np.pi / _lambda0
def __init__(self, shape, dperp, k, dz, npad=2):
super(ASMPropagator, self).__init__(shape, npad)
self._dperp = dperp
self._k = k
self._dz = dz
# calculate amount of padding
self._pad_width = (np.asarray(self._shape) * (self._npad - 1) / 2).astype(int)
self.kernel = self._init_kernel()
def _init_kernel(self):
# construct Rayleigh-Sommerfeld transfer function (Goodman: Fourier Optics, 4.2.3)
wl = self._k0 / self._k # relative wavelength
shape_pad = np.asarray(self._shape) + 2 * self._pad_width
f = fftfreqn(shape_pad, self._dperp) # spatial frequencies
f2 = squaresum(f)
mask = f2 < 1 / wl
phasechirp = np.sqrt(1 - np.square(wl) * (mask * f2))
kernel = mask * np.exp(1j * self._k * self._dz * phasechirp)
return kernel
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment