Commit 8dc38167 authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

split up fdsolver and propagation. introduce wavelength as parameter

parent b4693942
This diff is collapsed.
This diff is collapsed.
%% Cell type:code id: tags:
``` python
import sys
import numpy as np
import xraylib as xrl
import matplotlib.pyplot as plt
import fresnel.solver as solver
import fresnel.propagate as propagate
```
%% Cell type:code id: tags:
``` python
%matplotlib inline
```
%% Cell type:code id: tags:
``` python
energy = 12. # keV
# geometry
wavelength = 1.2398 / (energy * 1e3) * 1e-3 # mm
nx = 1024
ny = 512
nz = 1000
ztot = 0.4 # mm
wg_radius = 25e-6
xtot = 10 * wg_radius # mm
# units
dx = xtot / nx / wavelength
dz = ztot / nz / wavelength
print(f"dx: {dx}, dz: {dz} wavelengths")
```
%% Output
dx: 2.36303234392644, dz: 3871.592192289079 wavelengths
%% Cell type:code id: tags:
``` python
density_Ge = xrl.ElementDensity(xrl.SymbolToAtomicNumber('Ge'))
density_Polyimide = 1.42
density_Si = xrl.ElementDensity(xrl.SymbolToAtomicNumber('Si'))
n_Ge = xrl.Refractive_Index('Ge', energy, density_Ge)
n_Si = xrl.Refractive_Index('Si', energy, density_Si)
n_PI = xrl.Refractive_Index('Kapton Polyimide Film', energy, density_Polyimide)
xx = np.linspace(-xtot/2, xtot/2, nx)
core = (np.abs(xx) < wg_radius)
cladding = ~core
refractive_index = n_Ge * cladding + 1 * core
```
%% Cell type:code id: tags:
``` python
boundary = (0, 0,)
u0 = np.exp(-xx**2 / (2 * (2 * wg_radius)**2)).astype(np.complex128)
```
%% Cell type:code id: tags:
``` python
propagator = solver.PropagatorCS(refractive_index, u0, dz, dx)
propagator = propagate.FDPropagatorCS(refractive_index, u0, dz, dx)
```
%% Cell type:code id: tags:
``` python
field = np.zeros((nz, nx), dtype=np.complex128)
field[0,...] = u0
for iz in range(1, nz):
boundary = (0,0)
field[iz, ...] = propagator.step(refractive_index, boundary)
```
%% Cell type:code id: tags:
``` python
result = field.transpose()
fig, ax = plt.subplots()
im = ax.imshow(np.abs(result[:,:])**2, aspect='auto', clim=(0,3), cmap='jet', extent=[0, ztot, -xtot/2 *1e6, xtot/2 * 1e6])
ax.set_ylim(-2e6 * wg_radius, 2e6 * wg_radius )
fig.colorbar(im)
ax.set_xlabel(r'$z$ (mm)')
ax.set_ylabel(r'$x$ (nm)')
```
%% Output
Text(0, 0.5, '$x$ (nm)')
%% Cell type:code id: tags:
``` python
```
......
......@@ -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
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)]
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'):
......
......@@ -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
......
Supports Markdown
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