Commit 29cffa77 authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

Simplify0

parent 5dfa6715
......@@ -6,12 +6,12 @@ from . import _solver
class Solver2d:
"""
Solver for 2d PDEs of the form
Az(x) * u_z = -Axx(x) * u_xx + F(x,z) * u
A * u_z = -u_xx + F(x,z) * u
"""
dtype = np.complex128
def __init__(self, Az, Axx, F0, u0, dz, dx):
def __init__(self, A, F0, u0, dz, dx):
self._nx = u0.shape[-1]
self._ones = np.ones((self._nx,), dtype=self.dtype)
......@@ -22,18 +22,14 @@ class Solver2d:
self.dx = dx
self.dz = dz
self.rz = self._compute_rz(Az)
self.rxx = self._compute_rxx(Axx)
self.a = self._compute_a(A)
self.f = self._compute_f(F0)
def _compute_rz(self, Az):
return 2 * Az / self.dz * (1 + 0j)
def _compute_rxx(self, Axx):
return -Axx / self.dx ** 2 * (1 + 0j)
def _compute_a(self, A):
return A * self.dz / self.dx ** 2 * (1 + 0j)
def _compute_f(self, F):
return F * self._ones
return F * self.dz * self._ones
def step(self, F, boundary):
......@@ -48,7 +44,7 @@ class Solver2d:
u[0] = boundary[0]
u[-1] = boundary[1]
self.u = _solver.step1d_A0F(self.rz, self.rxx, fp, self.f, up, u)
self.u = _solver.step1d_A0F(self.a, fp, self.f, up, u)
return self.u
......@@ -56,12 +52,12 @@ class Solver2d:
class Solver2dSym:
"""
Solver for 2d PDEs of the form
Az * u_z = Axx * u_xx + Ax / x * u_x + F(x,z) * u
A * u_z = -u_xx - 1 / x * u_x + F(x,z) * u
"""
dtype = np.complex128
def __init__(self, Az, Axx, Ax, F0, u0, dz, dx):
def __init__(self, A, F0, u0, dz, dx):
self._nx = u0.shape[-1]
self._ones = np.ones((self._nx,), dtype=self.dtype)
......@@ -72,22 +68,14 @@ class Solver2dSym:
self.dx = dx
self.dz = dz
self.rz = self._compute_rz(Az)
self.rxx = self._compute_rxx(Axx)
self.rx = self._compute_rx(Ax)
self.a = self._compute_a(A)
self.f = self._compute_f(F0)
def _compute_rz(self, Az):
return 2 * Az / self.dz * (1 + 0j)
def _compute_rxx(self, Axx):
return -Axx / self.dx ** 2 * (1 + 0j)
def _compute_rx(self, Ax):
return -Ax / self.dx ** 2 * (1 + 0j)
def _compute_a(self, A):
return A * self.dz / self.dx ** 2 * (1 + 0j)
def _compute_f(self, F):
return F * self._ones
return F * self.dz * self._ones
def step(self, F, boundary):
......@@ -101,7 +89,7 @@ class Solver2dSym:
# set boundary values
u[-1] = boundary[0]
self.u = _solver.step1d_AAF_sym(self.rz, self.rxx, self.rx, fp, self.f, up, u)
self.u = _solver.step1d_AAF_sym(self.a, fp, self.f, up, u)
return self.u
......@@ -116,7 +104,7 @@ class Solver3d:
dtype = np.complex128
def __init__(self, Az, Axx, Ayy, F0, u0, dz, dy, dx):
def __init__(self, A, F0, u0, dz, dy, dx):
self._nx = u0.shape[-1]
self._ny = u0.shape[-2]
......@@ -129,22 +117,19 @@ class Solver3d:
self.dy = dy
self.dz = dz
self.rz = self._compute_rz(Az)
self.rxx = self._compute_rxx(Axx)
self.ryy = self._compute_ryy(Ayy)
self.f = self._compute_f(F0)
self.ax = self._compute_ax(A)
self.ay = self._compute_ay(A)
def _compute_rz(self, Az):
return 2 * Az * 2 / self.dz * (1 + 0j)
self.f = self._compute_f(F0)
def _compute_rxx(self, Axx):
return -Axx / self.dx ** 2 * (1 + 0j)
def _compute_ax(self, A):
return A * self.dz / self.dx ** 2 * (1 + 0j)
def _compute_ryy(self, Ayy):
return -Ayy / self.dy ** 2 * (1 + 0j)
def _compute_ay(self, A):
return A * self.dz / self.dy ** 2 * (1 + 0j)
def _compute_f(self, F):
return F / 2 * self._ones
return F * self.dz * self._ones
def step(self, F, boundary):
......@@ -162,7 +147,7 @@ class Solver3d:
u[:, -1] = boundary[3]
# the C++ functions use Eigen3 internally, which requires C-Major data storage
u_T = _solver.step2d_A0F(self.rz, self.rxx, self.ryy, fp.T, self.f.T, up.T, u.T)
u_T = _solver.step2d_A0F(self.ax, self.ay, fp.T, self.f.T, up.T, u.T)
self.u = u_T.T
return self.u
import numpy as np
from functools import reduce
from scipy.fft import fftn, ifftn, ifftshift, fftshift
from . import finite_differences as fd
from .misc import fftfreqn
from . import vacuum
_lambda0 = 1.0 # all lengths are in units of the wavelength
_k0 = 2 * np.pi / _lambda0
def _compute_potential(n_, k_=_k0):
return k_ * k_ * (1 - n_ * n_)
def squaresum(a):
return reduce(np.add, map(np.square, a))
def rayleighSommerfeldTF(shape, dperp, k, dz):
# construct Rayleigh-Sommerfeld transfer function (Goodman: Fourier Optics, 4.2.3)
wl = _k0 / k # relative wavelength
f = fftfreqn(shape, dperp) # spatial frequencies
f2 = squaresum(f)
mask = f2 < 1 / wl
phasechirp = np.sqrt(1 - wl ** 2 * (mask * f2))
# phasechirp = np.sqrt(1 - wl ** 2 * f2)
TF = mask * np.exp(1j * k * dz * phasechirp)
# TF = np.exp(1j * k * dz * phasechirp)
return TF
return -1j * k_ / 2 * (1 - n_ * n_)
class MultislicePropagator:
......@@ -60,37 +36,32 @@ class MultislicePropagator:
self._ones = np.ones(u0.shape, dtype=self.dtype)
self._k = _k0 / wl
self._fourierKernel = rayleighSommerfeldTF(self._ones.shape, self._dperp, self._k, self._dz)
# ASM Propagator without padding
self._fs_propagator = vacuum.ASMPropagator(
self._ones.shape, self._dperp, self._dz, wl, npad=1
)
self.u = ifftshift(u0)
self.u = u0
def step(self, n, workers=-1):
nshift = ifftshift(n)
F = _compute_potential(n, self._k) * self._ones
F = _compute_potential(nshift, self._k) * self._ones
real_kernel = np.exp(self._dz * F)
up = real_kernel * self.u
self.u = self._fs_propagator(up, workers)
realKernel = np.exp(-1j * self._dz / (2 * self._k) * F)
up = self.u
up *= realKernel
ft_u = fftn(up, workers=workers)
ft_u *= self._fourierKernel
self.u = ifftn(ft_u, workers=workers)
return fftshift(self.u)
return self.u
class FDPropagator2d(fd.Solver2d):
def __init__(self, n0, u0, dz, dx, wl=1.0):
self._k = _k0 / wl
Az = 2j * self._k
Axx = 1
A = 1j / (2 * self._k)
F0 = _compute_potential(n0, self._k)
super().__init__(Az, Axx, F0, u0, dz, dx)
super().__init__(A, F0, u0, dz, dx)
def step(self, n, boundary):
......@@ -103,13 +74,10 @@ class FDPropagatorCS(fd.Solver2dSym):
def __init__(self, n0, u0, dz, dx, wl=1.0):
self._k = _k0 / wl
Az = 2j * self._k
Axx = 1
Ax = 1 # 1/x-dependency is handled in C++-code
A = 1j / (2 * self._k)
F0 = _compute_potential(n0, self._k)
super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
super().__init__(A, F0, u0, dz, dx)
def step(self, n, boundary):
......@@ -122,13 +90,10 @@ class FDPropagator3d(fd.Solver3d):
def __init__(self, n0, u0, dz, dy, dx, wl=1.0):
self._k = _k0 / wl
Az = 2j * self._k
Axx = 1
Ayy = 1
A = 1j / (2 * self._k)
F0 = _compute_potential(n0, self._k)
super().__init__(Az, Axx, Ayy, F0, u0, dz, dy, dx)
super().__init__(A, F0, u0, dz, dy, dx)
def step(self, n, boundary):
......
......@@ -253,15 +253,7 @@ class ASMPropagator(FTConvolutionPropagator):
mask = f2 < 1 / np.square(self._wl) if self.mask_evanescent else 1
# expand dimension of dz
dz = self._dz.reshape(
[
-1,
]
+ self._ndim
* [
1,
]
)
dz = self._dz.reshape([-1] + self._ndim * [1])
kernel = mask * np.exp(1j * self._k0 * dz * phase_chirp[np.newaxis, ...])
......
......@@ -73,8 +73,7 @@ namespace finite_differences {
*/
const array_1D step1d_A0F(const scalar rz,
const scalar rxx,
const array_1D step1d_A0F(const scalar a,
const Eigen::Ref<const array_1D> fp,
const Eigen::Ref<const array_1D> f,
const Eigen::Ref<const array_1D> up,
......@@ -92,16 +91,15 @@ namespace finite_differences {
// Note that matrix indices go from 1 to u.size()-1, so that indices are off by 1
for (size_t i = 1; i <= n; ++i) {
m(0, i-1) = - rxx; // lower
m(1, i-1) = rz + 2.0 * rxx - f(i); // diagonal
m(2, i-1) = - rxx; // upper;
r(i-1) = (rz + fp(i)) * up(i)
+ rxx * (up(i-1) - 2.0 * up(i) + up(i+1));
m(0, i-1) = - a; // lower
m(1, i-1) = 2.0 + 2.0 * a - f(i); // diagonal
m(2, i-1) = - a; // upper;
r(i-1) = a * up(i-1) + (2.0 - 2.0 * a + fp(i)) * up(i) + a * up(i+1);
}
// boundary
r(0) += rxx * u(0);
r(n-1) += rxx * u(n+1);
r(0) += a * u(0);
r(n-1) += a * u(n+1);
// initialize writeable copy
array_1D u_new {u};
......@@ -111,9 +109,7 @@ namespace finite_differences {
}
const array_1D step1d_AAF_sym(const complex rz,
const complex rxx,
const complex rx,
const array_1D step1d_AAF_sym(const complex a,
const Eigen::Ref<const array_1D> fp,
const Eigen::Ref<const array_1D> f,
const Eigen::Ref<const array_1D> up,
......@@ -129,27 +125,26 @@ namespace finite_differences {
array_1D r = array_1D::Zero(n, 1);
for (size_t i = 1; i < n; ++i) {
complex rhix = rx / static_cast<complex>(i) / 2.0;
m(0, i) = -rxx + rhix; // lower
m(1, i) = rz + 2.0 * rxx - f(i); // diagonal
m(2, i) = -rxx - rhix; // upper;
r(i) = (rz + f(i)) * up(i)
+ rxx * (up(i-1) - 2.0 * up(i) + up(i+1))
+ rhix * (up(i+1) - up(i-1));
complex rhix = 1.0 / static_cast<complex>(i);
m(0, i) = -(1.0 - 0.5 * rhix) * a; // lower
m(1, i) = 2.0 + 2.0 * a - f(i); // diagonal
m(2, i) = -(1.0 + 0.5 * rhix) * a; // upper;
r(i) = (1.0 - 0.5 * rhix) * a * up(i-1)
+ (2.0 - 2.0 * a + fp(i)) * up(i)
+ (1.0 + 0.5 * rhix) * a * up(i+1);
}
// left boundary
m(0, 0) = 0; // lower
m(1, 0) = rz + 2.0 * rxx + 2.0 * rx - f(0); // diagonal
m(2, 0) = -2.0 * rxx - 2.0 * rx; // upper;
r(0) = (rz + f(0)) * up(0)
+ (2.0 * rxx + 2.0 * rx) * (up(1) - up(0));
m(1, 0) = 2.0 + 4.0 * a - f(0); // diagonal
m(2, 0) = -4.0 * a; // upper;
r(0) = (2.0 - 4.0 * a + f(0)) * up(0) + 4.0 * a * up(1);
// right boundary
complex rxnh = rx / static_cast<complex>(n-1) / 2.0;
r(n-1) += (rxx + rxnh) * u(n);
complex rxnh = 1.0 / static_cast<complex>(n-1);
r(n-1) += (1.0 + 0.5 * rxnh) * a * u(n);
// initialize writeable copy
array_1D u_new {u};
......@@ -163,19 +158,16 @@ namespace finite_differences {
/* Two-step alternating-direction Peaceman-Rachford scheme
(J. W. Thomas: Numerical Partial Differential Equations)
*/
const array_2D step2d_A0F(
const complex rz,
const complex rxx,
const complex ryy,
const complex ax,
const complex ay,
const Eigen::Ref<const array_2D> &fp,
const Eigen::Ref<const array_2D> &f,
const Eigen::Ref<const array_2D> &up,
const Eigen::Ref<const array_2D> &u) {
// rxx and ryy differ by a factor of 1/2 from the definitions in Fuhse et al. (2005)
// .p : i
// . : i+1
// TODO(Leon): interpolate for half-step?
// coordinate system:
// numpy (z,y,x), eigen (rows, cols): x <-> cols, y <-> rows
......@@ -195,18 +187,16 @@ namespace finite_differences {
array_1D r = array_1D::Zero(nx, 1);
for (size_t ix = 1; ix <= nx; ++ix) {
m(0, ix-1) = - rxx;
m(1, ix-1) = rz + 2.0 * rxx - f(ix, iy);
m(2, ix-1) = - rxx;
m(0, ix-1) = -ax;
m(1, ix-1) = 2.0 + 2.0 * ax - 0.5 * f(ix, iy);
m(2, ix-1) = -ax;
r(ix-1) = (rz + fp(ix, iy)) * up(ix, iy)
+ ryy * (up(ix, iy - 1) - 2.0 * up(ix, iy) + up(ix, iy + 1));
r(ix-1) = ay * up(ix, iy-1) + (2.0 - 2.0 * ay + 0.5 * fp(ix, iy)) * up(ix, iy) + ay * up(ix, iy+1);
}
// boundary conditions
// TODO(Leon): use interpolated values?
r(0) += rxx * u(0, iy);
r(nx-1) += rxx * u(nx+1, iy);
r(0) += ax * u(0, iy);
r(nx-1) += ax * u(nx+1, iy);
uhalf.col(iy).segment(1, nx) = algebra::tridiagonal(std::move(m), std::move(r));
}
......@@ -215,6 +205,8 @@ namespace finite_differences {
// second halfstep
// note: internal storage is column-major such that column-wise access would be much faster
// TODO(Leon): transpose up, fp, and f
array_2D u_new_trans = array_2D::Zero(u.cols(), u.rows()); // transposed shape
#pragma omp parallel for
......@@ -223,16 +215,16 @@ namespace finite_differences {
array_1D r = array_1D::Zero(ny, 1);
for (size_t iy = 1; iy <= ny; ++iy) {
m(0, iy-1) = - ryy;
m(1, iy-1) = rz + 2. * ryy - f(ix, iy);
m(2, iy-1) = - ryy;
m(0, iy-1) = -ay;
m(1, iy-1) = 2.0 + 2.0 * ay - 0.5 * f(ix, iy);
m(2, iy-1) = -ay;
r(iy-1) = (rz + fp(ix, iy)) * uhalf(ix, iy)
+ rxx * (uhalf(ix-1, iy) - 2.0 * uhalf(ix, iy) + uhalf(ix+1, iy));
r(iy-1) = ax * uhalf(ix-1, iy) + (2.0 - 2.0 * ax + 0.5 * fp(ix, iy)) * uhalf(ix, iy) + ax * uhalf(ix+1, iy);
}
r(0) += ryy * u(ix, 0);
r(ny-1) += ryy * u(ix, ny+1);
// boundary conditions
r(0) += ay * u(ix, 0);
r(ny-1) += ay * u(ix, ny+1);
u_new_trans.col(ix).segment(1, ny) = algebra::tridiagonal(std::move(m), std::move(r));
}
......
......@@ -9,8 +9,7 @@ namespace fd = finite_differences;
PYBIND11_MODULE(_solver, m) {
m.def("step1d_A0F",
&fd::step1d_A0F,
py::arg("rz").noconvert(),
py::arg("rxx").noconvert(),
py::arg("a").noconvert(),
py::arg("fp").noconvert(),
py::arg("f").noconvert(),
py::arg("up").noconvert(),
......@@ -19,9 +18,8 @@ PYBIND11_MODULE(_solver, m) {
m.def("step2d_A0F",
&fd::step2d_A0F,
py::arg("rz").noconvert(),
py::arg("rxx").noconvert(),
py::arg("ryy").noconvert(),
py::arg("ax").noconvert(),
py::arg("ay").noconvert(),
py::arg("fp").noconvert(),
py::arg("f").noconvert(),
py::arg("up").noconvert(),
......@@ -30,9 +28,7 @@ PYBIND11_MODULE(_solver, m) {
m.def("step1d_AAF_sym",
&fd::step1d_AAF_sym,
py::arg("rz").noconvert(),
py::arg("rxx").noconvert(),
py::arg("rx").noconvert(),
py::arg("a").noconvert(),
py::arg("fp").noconvert(),
py::arg("f").noconvert(),
py::arg("up").noconvert(),
......
Markdown is supported
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