split up fdsolver and propagation. introduce wavelength as parameter

examples/TestCS.ipynb

@@ -6,12 +6,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sys\n",
     "import numpy as np\n",
     "import xraylib as xrl\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "import fresnel.solver as solver"
+    "import fresnel.propagate as propagate"
    ]
   },
   {
@@ -96,7 +95,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "propagator = solver.PropagatorCS(refractive_index, u0, dz, dx)"
+    "propagator = propagate.FDPropagatorCS(refractive_index, u0, dz, dx)"
    ]
   },
   {

fresnel/solver.py → fresnel/finite_differences.py

@@ -2,8 +2,6 @@ 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():
     """
@@ -186,63 +184,4 @@ class Solver3d():
         u_T = _solver.step2d_A0Fs(self.rz, self.rxx, self.ryy, fp.T, self.f.T, up.T, u.T)
         self.u = u_T.T
 
-        return self.u
-
-
-class Propagator2d(Solver2d):
-
-    def __init__(self, n0, u0, dz, dx):        
-
-        Az = 2j * _k
-        Axx = 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)
-        
-        Az = 2j * _k * self._x
-        Axx = self._x
-        Ax  = 1
-        F0 = _compute_potential(n0)
-
-        super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
-
-
-    def step(self, n, boundary):
-
-        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 = _compute_potential(n)
-
-        return super().step(F, boundary)
+        return self.u
\ No newline at end of file

fresnel/misc.py

+import numpy as np
+from pyfftw.interfaces.numpy_fft import fftfreq, ifftshift, fftshift
+
+
+def fftfreqn(N, d=1.0):
+    ndim = len(N)
+    
+    d *= np.ones(ndim)
+
+    # 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)]
+
+    return xi
+
+
+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)]

fresnel/propagation.py → fresnel/propagate.py

-from pyfftw.interfaces.numpy_fft import fftn,fftfreq,ifftn,ifftshift,fftshift
+from pyfftw.interfaces.numpy_fft import fft,fftn,fftfreq,ifftn,ifftshift,fftshift
 import numpy as np
+from . import finite_differences as fd
+from .misc import fftfreqn, crop
 
-_k = 2 * np.pi # all lengths are in units of the wavelength
-_compute_potential = lambda n : _k**2 * (1 - n*n)
+_lambda0 = 1.  # all lengths are in units of the wavelength
+_k0 = 2 * np.pi / _lambda0
+_compute_potential = lambda n_, k_=_k0 : k_*k_ * (1 - n_*n_)
 
-def fftfreqn(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')
+def fresnelKernel(shape, fresnelNumbers, method='TF'):
 
-    xi = [np.reshape(2 * np.pi * fftfreq(N[dim]), tuple(shapes[dim, :])) for dim in range(ndim)]
-
-    return xi
-
-
-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)
+    ndim = len(shape)
 
     if np.isscalar(fresnelNumbers):
         fresnelNumbers = fresnelNumbers * np.ones(ndim)
 
-    kernel = np.ones(N, dtype=np.complex128)
-    xi = fftfreqn(N)
+    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)
 
-    for dim in range(ndim):
-        kernel *= np.exp((-1j/(2*_k*fresnelNumbers[dim])) * xi[dim]**2)
+        return kernel
 
-    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 FresnelPropagator2d():
+    
+class FresnelPropagator():
     """
     Multi-Slice approximation
     Paganin, Coherent X-Ray Optics, p101
     """
 
     dtype = np.complex128
+    
+    
+    def __init__(self, u0, d, wl=1., method='TF'):
+        
+        ndim = len(d)
+        
+        # TODO: check input
+        # assert u0 is array
+        # assert d is tuple
+        # assert wl is scalar numeric
+        # assert ndim == u0.ndim + 1
 
-    def __init__(self, u0, dz, dx):
+        self._dz = d[0]
+        self._dperp = np.array(d[1:])
+        self._ones = np.ones(u0.shape, dtype = self.dtype)
+        self._k = _k0 / wl
 
-        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)
+        fresnelNumbers = self._dperp**2 / (self._dz * wl) 
+        self._fourierKernel = fresnelKernel(self._ones.shape, fresnelNumbers, method)
         
         self.u = ifftshift(u0)
         
-        
+    
     def step(self, n):
         
-        F = _compute_potential(ifftshift(n)) * self._ones
+        nshift = ifftshift(n)
         
-        self._realKernel = np.exp(-1j / (2 * _k) * self._dz * F )
+        F = _compute_potential(nshift, self._k) * self._ones
+        
+        self._realKernel = np.exp(-1j / (2 * self._k) * self._dz * F )
         up = self.u
         
         self.u = ifftn(self._fourierKernel * fftn(self._realKernel * up))
@@ -69,38 +88,71 @@ class FresnelPropagator2d():
         return fftshift(self.u)
     
 
-class FresnelPropagator3d():
-    """
-    Multi-Slice approximation
-    Paganin, Coherent X-Ray Optics, p101
-    """
-
-    dtype = np.complex128
+class FDPropagator2d(fd.Solver2d):
 
-    def __init__(self, u0, dz, dy, dx):
+    def __init__(self, n0, u0, dz, dx, wl=1.):        
 
-        self._nx = u0.shape[-1]
-        self._nx = u0.shape[-2]
-        self._ones = np.ones((self._ny, self._nx,), dtype = self.dtype)
+        self._k = _k0 / wl
 
-        self._dz = dz
-        
-        fresnelNumber = dx**2 / dz # in units of wavelength
-        self._fourierKernel = fresnelKernel(self._ones.shape, fresnelNumber)
-        
-        self.u = ifftshift(u0)
+        Az = 2j * self._k
+        Axx = 1
         
+        F0 = _compute_potential(n0, self._k)
+
+        super().__init__(Az, Axx, F0, u0, dz, dx)
+
+
+    def step(self, n, boundary):
+
+        F = _compute_potential(n, self._k)
+
+        return super().step(F, boundary)
+
+
+class FDPropagatorCS(fd.Solver2dfull):
+
+    def __init__(self, n0, u0, dz, dx, wl=1.):        
+
+        nx = u0.shape[-1]
+        self._x = np.linspace(-nx * dx / 2, nx * dx / 2, nx)
         
-    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))
+        self._k = _k0 / wl
+
+        Az = 2j * self._k * self._x
+        Axx = self._x
+        Ax  = 1
+        F0 = _compute_potential(n0, self._k) * self._x
+
+        super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
+
+
+    def step(self, n, boundary):
+
+        F = _compute_potential(n, self._k) * self._x
+
+        return super().step(F, boundary)
+    
+
+class FDPropagator3d(fd.Solver3d):
+
+    def __init__(self, n0, u0, dz, dy, dx, wl=1.):        
+   
+        self._k = _k0 / wl
+
+        Az = 2j * self._k
+        Axx = 1
+        Ayy = 1
+        F0 = _compute_potential(n0, self._k)
+
+        super().__init__(Az, Axx, Ayy, F0, u0, dz, dy, dx)
+
         
-        return fftshift(self.u)
+    def step(self, n, boundary):
+
+        F = _compute_potential(n, self._k)
+
+        return super().step(F, boundary)
+
 
 
 def fresnelPropagate(im, fresnelNumbers, sizePad=None, sizeOut=None, padMode='edge'):

source/finite_difference.h

@@ -2,7 +2,7 @@
 
 #include <complex>
 #include <Eigen/Dense>
-#include <iostream>
+//#include <iostream>
 
 namespace algebra{
 
@@ -174,7 +174,7 @@ namespace finite_differences{
     size_t nx = u.rows() - 2;
     size_t ny = u.cols() - 2;
       
-    std::cout << "nx " << u.cols() << ", ny " << u.rows() << std::endl; 
+//    std::cout << "nx " << u.cols() << ", ny " << u.rows() << std::endl; 
     
     // first halfstep
     // x at i + 1/2
@@ -264,7 +264,7 @@ namespace finite_differences{
     size_t nx = u.rows() - 2;
     size_t ny = u.cols() - 2;
       
-    std::cout << "nx " << u.cols() << ", ny " << u.rows() << std::endl; 
+//    std::cout << "nx " << u.cols() << ", ny " << u.rows() << std::endl; 
     
     // first halfstep
     // x at i + 1/2