Commit 97a5b558 authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

add discrete hankel transform and vacuum propagation

parent bc06658c
import numpy as np
import scipy.special
def hankelMatrix(N, n=0):
'''
returns a N x N matrix for discrete Hankel transfrom of n-th order.
N: number of pixels
n: order of Hankel transform
The Hankel matrix is self-inverse! I.e. HH = Id
For forward and backward Hankel transfrom different prefactors have to be considered!
As in: Theory and operational rules for the discreteHankel transform
by Natalie Baddour* and Ugo Chouinard
https://doi.org/10.1364/JOSAA.32.000611
'''
jn = np.array(scipy.special.jn_zeros(n, N + 1))
k = np.expand_dims(np.arange(N), axis=0)
m = np.expand_dims(np.arange(N), axis=1)
jN = jn[-1]
Y = scipy.special.jn(n, jn[m] * jn[k] / jN) # matrix
Y *= 2 / (jN * scipy.special.jn(n + 1, jn[k]) ** 2) # prefactor
return Y
def hankelFreq(N, n=0, kmax=0.5):
'''
Returns the Hankel space (frequency) sampling grid for the inverse discrete
Hankel transfrom (of order n) of a signal with N pixels.
kmax is the maximum sampling frequency in dimensionless units, i.e.
minimal sampled realspace oscillation 2px -> max. sampled frequency 1/(2px)
-> 0.5 dimensionless
'''
jn = np.array(scipy.special.jn_zeros(n, N + 1))
return jn[:-1] * kmax / jn[N]
def hankelSamples(N, n=0, kmax=0.5):
'''
Returns the real space sampling grid for the forward discrete Hankel
transfrom (of order n) of a signal with N pixels.
kmax is the maximum sampling frequency in dimensionless units, i.e.
minimal sampled realspace oscillation 2px -> max. sampled frequency 1/(2px)
-> 0.5 dimensionless
'''
jn = np.array(scipy.special.jn_zeros(n, N))
return jn / (kmax*2*np.pi)
class DiscreteHankelTransform:
def __init__(self, N, n=0, kmax=0.5):
self._matrix = hankelMatrix(N, n)
def __call__(self, x):
return self._matrix @ x
......@@ -140,71 +140,3 @@ class FDPropagator3d(fd.Solver3d):
return super().step(F, boundary)
def fresnelKernel(shape, fresnelNumbers, method='TF'):
ndim = len(shape)
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(ndim)
kernel = np.ones(shape, dtype=np.complex128)
if method is 'TF':
# transfer function (TF) method
xi = fftfreqn(shape, 1)
for dim in range(ndim):
kernel *= np.exp((-1j/(2*_k0*fresnelNumbers[dim])) * xi[dim]**2)
return kernel
if method is 'IR':
# impulse response (IR) method
# phase factor
kernel *= np.prod(np.sqrt(abs(fresnelNumbers))) * np.exp((-1j*np.pi/4) * np.sign(fresnelNumbers))
# TODO: use dedicated function for this to avoid loss of precision
x = fftfreqn(shape, 1/shape)
for dim in range(ndim):
kernel *= fft(np.exp(1j*np.pi*fresnelNumbers[dim] * x[dim]**2))
return kernel
raise ValueError(f'Invalid method "{method}". Must be "TF" or "IR".')
def fresnelPropagate(im, fresnelNumbers, sizePad=None, sizeOut=None, padMode='edge'):
sizeIn = np.array(im.shape)
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(im.ndim)
Fpix = fresnelNumbers / sizeIn**2
if sizeOut is None:
sizeOut = sizeIn
if sizePad is None:
sizePad = sizeIn + 2 * np.ceil(0.5 / abs(Fpix)).astype('int')
sizePad = np.maximum(sizePad, sizeOut)
if np.any(sizePad >1.2 * (sizeIn + sizeOut)) and np.prod(sizePad) > 4096**2:
raise OverflowError('Padded size {} excessively large. Aborting.'.format(sizePad))
padPost = (sizePad - sizeIn)//2
padPre = sizePad - sizeIn - padPost
cropPost = (sizePad - sizeOut)//2
cropPre = sizePad - sizeOut - cropPost
pad_width = list(zip(padPre, padPost))
crop_width = list(zip(cropPre, cropPost))
imPad = np.pad(im, pad_width, mode='edge')
propKernel = propKernFresnel(sizePad, Fpix)
imProp = ifftn( propKernel * fftn(imPad))
return crop(imProp, crop_width)
from . import hankel
import numpy as np
from pyfftw.interfaces.numpy_fft import fft,fftn,fftfreq,ifftn,ifftshift,fftshift
_lambda0 = 1. # all lengths are in units of the wavelength
_k0 = 2 * np.pi / _lambda0
def fresnelKernelCS(N, fresnelNumber):
hFreq = hankel.hankelFreq(N)
kern = np.exp(-1j * np.pi * hFreq**2 / fresnelNumber )
return kern
class FresnelPropagatorCS:
def __init__(self, N, fresnelNumber):
self._N = N
Y = hankel.hankelMatrix(N, 0)
kern = fresnelKernelCS(N, fresnelNumber)
self._matrix = Y @ (kern @ Y)
def __call__(u):
uprop = self._matrix @ (u - 1) + 1
return uprop
def fresnelKernel(shape, fresnelNumbers, method='TF'):
ndim = len(shape)
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(ndim)
kernel = np.ones(shape, dtype=np.complex128)
if method is 'TF':
# transfer function (TF) method
xi = fftfreqn(shape, 1)
for dim in range(ndim):
kernel *= np.exp((-1j/(2*_k0*fresnelNumbers[dim])) * xi[dim]**2)
return kernel
if method is 'IR':
# impulse response (IR) method
# phase factor
kernel *= np.prod(np.sqrt(abs(fresnelNumbers))) * np.exp((-1j*np.pi/4) * np.sign(fresnelNumbers))
# TODO: use dedicated function for this to avoid loss of precision
x = fftfreqn(shape, 1/shape)
for dim in range(ndim):
kernel *= fft(np.exp(1j*np.pi*fresnelNumbers[dim] * x[dim]**2))
return kernel
raise ValueError(f'Invalid method "{method}". Must be "TF" or "IR".')
class FresnelPropagator:
def __init__(self, shape, fresnelNumbers, method='TF'):
self._shape = np.array(shape)
self._ndim = len(self._shape)
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(self._ndim)
self._kernel = fresnelKernel(shape, fresnelNumbers, method)
def __call__(u):
uprop = ifftn( self._kernel * fftn(u))
"""
def fresnelPropagate(im, fresnelNumbers, sizePad=None, sizeOut=None, padMode='edge'):
sizeIn = np.array(im.shape)
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(im.ndim)
Fpix = fresnelNumbers / sizeIn**2
if sizeOut is None:
sizeOut = sizeIn
if sizePad is None:
sizePad = sizeIn + 2 * np.ceil(0.5 / abs(Fpix)).astype('int')
sizePad = np.maximum(sizePad, sizeOut)
if np.any(sizePad >1.2 * (sizeIn + sizeOut)) and np.prod(sizePad) > 4096**2:
raise OverflowError('Padded size {} excessively large. Aborting.'.format(sizePad))
padPost = (sizePad - sizeIn)//2
padPre = sizePad - sizeIn - padPost
cropPost = (sizePad - sizeOut)//2
cropPre = sizePad - sizeOut - cropPost
pad_width = list(zip(padPre, padPost))
crop_width = list(zip(cropPre, cropPost))
imPad = np.pad(im, pad_width, mode='edge')
propKernel = propKernFresnel(sizePad, Fpix)
imProp = ifftn( propKernel * fftn(imPad))
return crop(imProp, crop_width)
"""
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