implement fresnel propagator

fresnel/propagation.py

 from pyfftw.interfaces.numpy_fft import fftn,fftfreq,ifftn,ifftshift,fftshift
 import numpy as np
 
+_k = 2 * np.pi # all lengths are in units of the wavelength
+_compute_potential = lambda n : _k**2 * (1 - n*n)
+
 def fftfreqn(N):
-  ndim = len(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')
@@ -11,43 +14,112 @@ def fftfreqn(N):
 
     return xi
 
-def propKernFresnel(N, fresnelNumbers):
 
-  ndim = len(N)
+def crop(a, crop_width):
+    slices = [slice(w[0], -w[1] if w[1] > 0 else None) for w in crop_width]
+
+    return a[tuple(slices)]
+
+
+def fresnelKernel(N, fresnelNumbers):
+
+    ndim = len(N)
 
     if np.isscalar(fresnelNumbers):
-      fresnelNumbers = fresnelNumbers * np.ones(ndim)
+        fresnelNumbers = fresnelNumbers * np.ones(ndim)
 
     kernel = np.ones(N, dtype=np.complex128)
     xi = fftfreqn(N)
 
     for dim in range(ndim):
-      kernel *= np.exp((-1j/(4*np.pi*fresnelNumbers[dim])) * xi[dim]**2)
+        kernel *= np.exp((-1j/(2*_k*fresnelNumbers[dim])) * xi[dim]**2)
 
     return kernel
 
-def crop(a, crop_width):
-  slices = [slice(w[0], -w[1] if w[1] > 0 else None) for w in crop_width]
 
-    return a[tuple(slices)]
+class FresnelPropagator2d():
+    """
+    Multi-Slice approximation
+    Paganin, Coherent X-Ray Optics, p101
+    """
+
+    dtype = np.complex128
+
+    def __init__(self, u0, dz, dx):
+
+        self._nx = u0.shape[-1]
+        self._ones = np.ones((self._nx,), dtype = self.dtype)
+        self._dz = dz
+        
+        fresnelNumber = dx**2 / dz # in units of wavelength
+        self._fourierKernel = fresnelKernel(self._ones.shape, fresnelNumber)
+        
+        self.u = ifftshift(u0)
+        
+        
+    def step(self, n):
+        
+        F = _compute_potential(ifftshift(n)) * self._ones
+        
+        self._realKernel = np.exp(-1j / (2 * _k) * self._dz * F )
+        up = self.u
+        
+        self.u = ifftn(self._fourierKernel * fftn(self._realKernel * up))
+        
+        return fftshift(self.u)
+    
+
+class FresnelPropagator3d():
+    """
+    Multi-Slice approximation
+    Paganin, Coherent X-Ray Optics, p101
+    """
+
+    dtype = np.complex128
+
+    def __init__(self, u0, dz, dy, dx):
+
+        self._nx = u0.shape[-1]
+        self._nx = u0.shape[-2]
+        self._ones = np.ones((self._ny, self._nx,), dtype = self.dtype)
+
+        self._dz = dz
+        
+        fresnelNumber = dx**2 / dz # in units of wavelength
+        self._fourierKernel = fresnelKernel(self._ones.shape, fresnelNumber)
+        
+        self.u = ifftshift(u0)
+        
+        
+    def step(self, n):
+        
+        F = _compute_potential(ifftshift(n)) * self._ones
+        
+        self._realKernel = np.exp(-1j / (2 * _k) * self._dz * F )
+        up = self.u
+        
+        self.u = ifftn(self._fourierKernel * fftn(self._realKernel * up))
+        
+        return fftshift(self.u)
+
 
 def fresnelPropagate(im, fresnelNumbers, sizePad=None, sizeOut=None, padMode='edge'):
-  sizeIn = np.array(im.shape)
+    sizeIn = np.array(im.shape)
 
     if np.isscalar(fresnelNumbers):
-      fresnelNumbers = fresnelNumbers * np.ones(im.ndim)
+        fresnelNumbers = fresnelNumbers * np.ones(im.ndim)
 
     Fpix = fresnelNumbers / sizeIn**2
 
     if sizeOut is None:
-      sizeOut = sizeIn
+        sizeOut = sizeIn
 
     if sizePad is None:
-      sizePad = sizeIn + 2 * np.ceil(0.5 / abs(Fpix)).astype('int')
+        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))
+        raise OverflowError('Padded size {} excessively large. Aborting.'.format(sizePad))
 
     padPost = (sizePad - sizeIn)//2
     padPre = sizePad - sizeIn - padPost
@@ -57,7 +129,6 @@ def fresnelPropagate(im, fresnelNumbers, sizePad=None, sizeOut=None, padMode='ed
     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))

fresnel/solver.py

@@ -3,6 +3,7 @@ import numpy as np
 from . import _solver
 
 _k = 2 * np.pi # all lengths are in units of the wavelength
+_compute_potential = lambda n : _k**2 * (1 - n*n)
 
 class Solver2d():
     """
@@ -12,6 +13,7 @@ class Solver2d():
 
     dtype = np.complex128
 
+    
     def __init__(self, Az, Axx, F0, u0, dz, dx):
 
         self._nx = u0.shape[-1]
@@ -58,41 +60,16 @@ class Solver2d():
         return self.u
 
 
-class Propagator2d(Solver2d):
-
-
-    def __init__(self, n0, u0, dz, dx):        
-
-        
-        Az = 2j * _k
-        Axx = 1
-        
-        F0 = self._compute_potential(n0)
-
-        super().__init__(Az, Axx, F0, u0, dz, dx)
-
-
-    def _compute_potential(self,n):
-
-        return _k**2 * (1 - n*n)
-
-
-    def step(self, n, boundary):
-
-        F = self._compute_potential(n)
-
-        return super().step(F, boundary)
-
 
 class Solver2dfull():
     """
     Solver for 2d PDEs of the form
       Az(x) * u_z = Axx(x) * u_xx + Ax(x) * u_x + F(x,z) * u
     """
-
  
     dtype = np.complex128
 
+    
     def __init__(self, Az, Axx, Ax, F0, u0, dz, dx):
 
         self._nx = u0.shape[-1]
@@ -142,42 +119,14 @@ class Solver2dfull():
         self.u = _solver.step1d_AAF(self.rz, self.rxx, self.rx, fp, self.f, up, u)
 
         return self.u
-
-
-
-class PropagatorCS(Solver2dfull):
-
-
-    def __init__(self, n0, u0, dz, dx):        
-
-        nx = u0.shape[-1]
-        self._x = np.linspace(-nx * dx / 2, nx * dx / 2, nx)
-        
-        Az = 2j * _k * self._x
-        Axx = self._x
-        Ax  = 1
-        F0 = self._compute_potential(n0)
-
-        super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
-
-
-    def _compute_potential(self,n):
-        
-        return _k**2 * (1 - n*n)
-
-
-
-    def step(self, n, boundary):
-
-        F = self._compute_potential(n) * self._x
-
-        return super().step(F, boundary)
     
 
 class Solver3d():
     """
     Solver for equations of the form
-      u_z = A * u_xx + A * u_yy + C * u
+     Az * u_z = Axx * u_xx + Ayy * u_yy + F(x,y,z) * u
+     
+    Az, Axx, Ayy are complex constants.
     """
 
     dtype = np.complex128
@@ -218,7 +167,6 @@ class Solver3d():
         return F * self.dz / 4. * self._ones
 
 
-
     def step(self, F, boundary):
 
         up = self.u
@@ -241,26 +189,60 @@ class Solver3d():
         return self.u
 
 
-class Propagator3d(Solver3d):
+class Propagator2d(Solver2d):
+
+    def __init__(self, n0, u0, dz, dx):        
 
-    def __init__(self, n0, u0, dz, dy, dx):        
-   
         Az = 2j * _k
         Axx = 1
-        Ayy = 1
+        F0 = _compute_potential(n0)
+
+        super().__init__(Az, Axx, F0, u0, dz, dx)
+
+
+    def step(self, n, boundary):
+
+        F = _compute_potential(n)
+
+        return super().step(F, boundary)
+
+
+class PropagatorCS(Solver2dfull):
+
+    def __init__(self, n0, u0, dz, dx):        
+
+        nx = u0.shape[-1]
+        self._x = np.linspace(-nx * dx / 2, nx * dx / 2, nx)
         
-        F0 = self._compute_potential(n0)
+        Az = 2j * _k * self._x
+        Axx = self._x
+        Ax  = 1
+        F0 = _compute_potential(n0)
 
-        super().__init__(Az, Axx, Ayy, F0, u0, dz, dy, dx)
+        super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
 
 
-    def _compute_potential(self,n):
+    def step(self, n, boundary):
 
-        return _k**2 * (1 - n*n)
+        F = _compute_potential(n) * self._x
 
+        return super().step(F, boundary)
+    
 
+class Propagator3d(Solver3d):
+
+    def __init__(self, n0, u0, dz, dy, dx):        
+   
+        Az = 2j * _k
+        Axx = 1
+        Ayy = 1
+        F0 = _compute_potential(n0)
+
+        super().__init__(Az, Axx, Ayy, F0, u0, dz, dy, dx)
+
+        
     def step(self, n, boundary):
 
-        F = self._compute_potential(n)
+        F = _compute_potential(n)
 
         return super().step(F, boundary)