Skip to content
GitLab
Commit 0b4461cd authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

improve vacuum propagation and add example

parent 97a5b558
Hide whitespace changes
Inline Side-by-side
examples/Vacuum.ipynb 0 → 100644
View file @ 0b4461cd
%% Cell type:code id: tags:
``` python
import fresnel.vacuum as vac
import numpy as np
import matplotlib.pyplot as plt
```
%% Cell type:code id: tags:
``` python
%matplotlib widget
```
%% Cell type:code id: tags:
``` python
n = 250
xx = np.linspace(-0.25, 0.25, n, endpoint=True).reshape(1,-1)
yy = np.linspace(-0.25, 0.25, n, endpoint=True).reshape(-1,1)
structuresize = 0.1
radius = structuresize/2
mask = (xx**2 < radius**2) * (yy**2 < radius**2)
u0 = mask * 1+0j
Fpix = 1/n
```
%% Cell type:code id: tags:
``` python
fig, ax = plt.subplots()
im = ax.imshow(np.abs(u0)**2)
fig.colorbar(im)
```
%% Output
<matplotlib.colorbar.Colorbar at 0x7fc8a219aa60>
%% Cell type:code id: tags:
``` python
# Voelz sampling criterion: wl * z / dx / L = wl * z / dx**2 / N
# pixel Fresnel Number = dx**2 / z / wl
# Fpix is chosen such that the Voelz criterion is satisfied!
# For other fresnel numbers, the sampling must be adapted (padding)
```
%% Cell type:code id: tags:
``` python
# Transfer function propagator
propTF = vac.FresnelTFPropagator(u0.shape, Fpix)
upropTF = propTF(u0)
fig, ax = plt.subplots()
ax.plot(xx.flatten(), np.abs(u0[n//2,:])**2)
ax.plot(xx.flatten(), np.abs(upropTF[n//2,:])**2)
```
%% Output
[<matplotlib.lines.Line2D at 0x7fc8a2101340>]
%% Cell type:code id: tags:
``` python
propIR = vac.FresnelIRPropagator(u0.shape, Fpix)
upropIR = propIR(u0)
fig, ax = plt.subplots()
ax.plot(xx.flatten(), np.abs(u0[n//2,:])**2)
ax.plot(xx.flatten(), np.abs(upropIR[n//2,:])**2)
```
%% Output
[<matplotlib.lines.Line2D at 0x7fc8a1cf65e0>]
%% Cell type:code id: tags:
``` python
mask2 = (xx**2 + yy**2 < radius**2)
u02 = mask2 * 1+0j
fig, ax = plt.subplots()
im = ax.imshow(np.abs(u02)**2)
fig.colorbar(im)
```
%% Output
<matplotlib.colorbar.Colorbar at 0x7fc8a1cda910>
%% Cell type:code id: tags:
``` python
propTF2 = vac.FresnelTFPropagator(u0.shape, 2.5 * Fpix)
u2propTF = propTF2(u02)
fig, ax = plt.subplots()
ax.plot(xx.flatten(), np.abs(u02[n//2,:])**2)
ax.plot(xx.flatten(), np.abs(u2propTF[n//2,:])**2)
```
%% Output
<ipython-input-49-f19e038b5538>:5: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).
fig, ax = plt.subplots()
[<matplotlib.lines.Line2D at 0x7fc8a1196a90>]
%% Cell type:code id: tags:
``` python
u02r = u02[n//2,:]
propCS = vac.FresnelPropagatorCS(u02r.shape[0], 2.5 * Fpix)
u2rprop = propCS(u02r)
```
%% Cell type:code id: tags:
``` python
fig, ax = plt.subplots()
ax.plot(xx.flatten(), np.abs(u02r)**2)
ax.plot(xx.flatten(), np.abs(u2rprop)**2)
```
%% Output
<ipython-input-51-29fa783d9e8a>:1: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).
fig, ax = plt.subplots()
[<matplotlib.lines.Line2D at 0x7fc8a10f64c0>]
fresnel/misc.py
View file @ 0b4461cd
......@@ -10,11 +10,22 @@ def fftfreqn(N, d=1.0):
# index trick: n'th line is a list of ones with a -1 on n'th position
shapes = np.ones((ndim, ndim), dtype='int') - 2 * np.eye(ndim, dtype='int')
xi = [np.reshape(2 * np.pi * fftfreq(N[dim], d[dim]), tuple(shapes[dim, :])) for dim in range(ndim)]
xi = [np.reshape(fftfreq(N[dim], d[dim]), tuple(shapes[dim, :])) for dim in range(ndim)]
return xi
def gridn(N):
ndim = len(N)
# index trick: n'th line is a list of ones with a -1 on n'th position
shapes = np.ones((ndim, ndim), dtype='int') - 2 * np.eye(ndim, dtype='int')
x = [np.reshape(np.arange(-N[dim]+1, N[dim]), tuple(shapes[dim, :])) for dim in range(ndim)]
return x
def crop(a, crop_width):
slices = [slice(w[0], -w[1] if w[1] > 0 else None) for w in crop_width]
......
fresnel/propagate.py
View file @ 0b4461cd
......@@ -3,7 +3,7 @@ import numpy as np
from functools import reduce
from . import finite_differences as fd
from .misc import fftfreqn, crop
from .misc import fftfreqn
_lambda0 = 1. # all lengths are in units of the wavelength
_k0 = 2 * np.pi / _lambda0
......@@ -19,7 +19,7 @@ def rayleighSommerfeldTF(shape, dperp, k, dz):
wl = _k0 / k # relative wavelength
f = fftfreqn(shape, dperp * 2*np.pi) # spatial frequencies
f = fftfreqn(shape, dperp) # spatial frequencies
f2 = squaresum(f)
mask = (f2 < 1 / wl)
......
fresnel/vacuum.py
View file @ 0b4461cd
from . import hankel
from .misc import fftfreqn, gridn
import numpy as np
from pyfftw.interfaces.numpy_fft import fft,fftn,fftfreq,ifftn,ifftshift,fftshift
from scipy.signal import fftconvolve
_lambda0 = 1. # all lengths are in units of the wavelength
_k0 = 2 * np.pi / _lambda0
......@@ -23,55 +26,35 @@ class FresnelPropagatorCS:
Y = hankel.hankelMatrix(N, 0)
kern = fresnelKernelCS(N, fresnelNumber)
self._matrix = Y @ (kern @ Y)
# TODO: express this nicer (kern is a diagonal matrix)
self._matrix = Y @ (kern[:,None] * Y)
def __call__(u):
def __call__(self, u):
uprop = self._matrix @ (u - 1) + 1
uprop = self._matrix @ u
return uprop
def fresnelKernel(shape, fresnelNumbers, method='TF'):
def fresnelTFKernel(shape, fresnelNumbers):
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)
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)
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".')
return kernel
class FresnelPropagator:
class FresnelTFPropagator:
def __init__(self, shape, fresnelNumbers, method='TF'):
def __init__(self, shape, fresnelNumbers):
self._shape = np.array(shape)
self._ndim = len(self._shape)
......@@ -79,42 +62,53 @@ class FresnelPropagator:
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(self._ndim)
self._kernel = fresnelKernel(shape, fresnelNumbers, method)
self._kernel = fresnelTFKernel(shape, fresnelNumbers)
def __call__(u):
def __call__(self, u):
uprop = ifftn( self._kernel * fftn(u))
"""
def fresnelPropagate(im, fresnelNumbers, sizePad=None, sizeOut=None, padMode='edge'):
sizeIn = np.array(im.shape)
return uprop
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(im.ndim)
Fpix = fresnelNumbers / sizeIn**2
if sizeOut is None:
sizeOut = sizeIn
def fresnelIRKernel(shape, fresnelNumbers):
if sizePad is None:
sizePad = sizeIn + 2 * np.ceil(0.5 / abs(Fpix)).astype('int')
sizePad = np.maximum(sizePad, sizeOut)
ndim = len(shape)
if np.any(sizePad >1.2 * (sizeIn + sizeOut)) and np.prod(sizePad) > 4096**2:
raise OverflowError('Padded size {} excessively large. Aborting.'.format(sizePad))
if np.isscalar(fresnelNumbers):
fresnelNumbers = fresnelNumbers * np.ones(ndim)
kernel = np.ones(2 * np.array(shape, dtype=int) - 1, dtype=np.complex128)
# phase factor
kernel *= np.prod(np.sqrt(np.abs(fresnelNumbers)) * np.exp((-1j * np.pi / 4) * np.sign(fresnelNumbers)) )
padPost = (sizePad - sizeIn)//2
padPre = sizePad - sizeIn - padPost
cropPost = (sizePad - sizeOut)//2
cropPre = sizePad - sizeOut - cropPost
x = gridn(shape)
for dim in range(ndim):
xdim = x[dim] # / shape[dim]
kernel *= np.exp(1j*np.pi*fresnelNumbers[dim] * xdim**2)
pad_width = list(zip(padPre, padPost))
crop_width = list(zip(cropPre, cropPost))
return kernel
class FresnelIRPropagator:
imPad = np.pad(im, pad_width, mode='edge')
propKernel = propKernFresnel(sizePad, Fpix)