improve vacuum propagation and add example

examples/Vacuum.ipynb

+%% 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: display_data
+
+
+%%%% Output: execute_result
+
+    <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: display_data
+
+
+%%%% Output: execute_result
+
+    [<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: display_data
+
+
+%%%% Output: execute_result
+
+    [<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: display_data
+
+
+%%%% Output: execute_result
+
+    <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: stream
+
+    <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()
+
+%%%% Output: display_data
+
+
+%%%% Output: execute_result
+
+    [<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: stream
+
+    <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()
+
+%%%% Output: display_data
+
+
+%%%% Output: execute_result
+
+    [<matplotlib.lines.Line2D at 0x7fc8a10f64c0>]

fresnel/misc.py

@@ -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

@@ -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

 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)
-    imProp = ifftn( propKernel * fftn(imPad))
+    def __init__(self, shape, fresnelNumbers):
+        
+        self._shape = np.array(shape)
+        self._ndim = len(self._shape)
+        
+        if np.isscalar(fresnelNumbers):
+            fresnelNumbers = fresnelNumbers * np.ones(self._ndim)
+        
+        self._kernel = fresnelIRKernel(shape, fresnelNumbers)
+    
+    
+    def __call__(self, u):
+        
+        uprop = fftconvolve(u, self._kernel, mode='valid')
 
-    return crop(imProp, crop_width)
-"""
+        return uprop
\ No newline at end of file