Commit 09b62083 authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

redo FD CS

parent d27e71fc
Pipeline #175056 failed with stages
in 39 seconds
This diff is collapsed.
This source diff could not be displayed because it is too large. You can view the blob instead.
This diff is collapsed.
......@@ -27,13 +27,13 @@ class Solver2d:
self.f = self._compute_f(F0)
def _compute_rz(self, Az):
return Az * self._ones
return Az * 2 * self.dx ** 2 / self.dz * (1 + 0j)
def _compute_rxx(self, Axx):
return -Axx * self.dz / 2.0 / self.dx ** 2 * self._ones
return -Axx * (1 + 0j)
def _compute_f(self, F):
return F * self.dz / 2.0 * self._ones
return F * self.dx ** 2 * self._ones
def step(self, F, boundary):
......@@ -53,10 +53,10 @@ class Solver2d:
return self.u
class Solver2dfull:
class Solver2dSym:
"""
Solver for 2d PDEs of the form
Az(x) * u_z = Axx(x) * u_xx + Ax(x) * u_x + F(x,z) * u
Az * u_z = Axx * u_xx + Ax / x * u_x + F(x,z) * u
"""
dtype = np.complex128
......@@ -78,16 +78,16 @@ class Solver2dfull:
self.f = self._compute_f(F0)
def _compute_rz(self, Az):
return Az * self._ones
return Az * 2 * self.dx ** 2 / self.dz * (1 + 0j)
def _compute_rxx(self, Axx):
return -Axx * self.dz / 2.0 / self.dx ** 2 * self._ones
return -Axx * (1 + 0j)
def _compute_rx(self, Ax):
return -Ax * self.dz / 4.0 / self.dx * self._ones
return -Ax * (1 + 0j)
def _compute_f(self, F):
return F * self.dz / 2.0 * self._ones
return F * self.dx ** 2 * self._ones
def step(self, F, boundary):
......@@ -99,10 +99,9 @@ class Solver2dfull:
u = self._boundary_slice
# set boundary values
u[0] = boundary[0]
u[-1] = boundary[1]
u[-1] = boundary[0]
self.u = _solver.step1d_AAF(self.rz, self.rxx, self.rx, fp, self.f, up, u)
self.u = _solver.step1d_AAF_sym(self.rz, self.rxx, self.rx, fp, self.f, up, u)
return self.u
......@@ -136,16 +135,16 @@ class Solver3d:
self.f = self._compute_f(F0)
def _compute_rz(self, Az):
return Az * (1 + 0j)
return Az * 2 * self.dx ** 2 / self.dz * (1 + 0j)
def _compute_rxx(self, Axx):
return -Axx * self.dz / 2.0 / self.dx ** 2 * (1 + 0j)
return -Axx * (1 + 0j)
def _compute_ryy(self, Ayy):
return -Ayy * self.dz / 2.0 / self.dy ** 2 * (1 + 0j)
return -Ayy * self.dx ** 2 / self.dy ** 2 * (1 + 0j)
def _compute_f(self, F):
return F * self.dz / 4.0 * self._ones
return F * self.dx ** 2 / 2 * self._ones
def step(self, F, boundary):
......@@ -163,7 +162,7 @@ class Solver3d:
u[:, -1] = boundary[3]
# the C++ functions use Eigen3 internally, which requires C-Major data storage
u_T = _solver.step2d_A0Fs(self.rz, self.rxx, self.ryy, fp.T, self.f.T, up.T, u.T)
u_T = _solver.step2d_A0F(self.rz, self.rxx, self.ryy, fp.T, self.f.T, up.T, u.T)
self.u = u_T.T
return self.u
......@@ -95,24 +95,23 @@ class FDPropagator2d(fd.Solver2d):
return super().step(F, boundary)
class FDPropagatorCS(fd.Solver2dfull):
class FDPropagatorCS(fd.Solver2dSym):
def __init__(self, n0, u0, dz, dx, wl=1.0):
nx = u0.shape[-1]
self._x = (np.arange(-nx//2, nx//2) + 0.5) * dx
self._k = _k0 / wl
Az = 2j * self._k * self._x
Axx = self._x
Ax = 1
F0 = _compute_potential(n0, self._k) * self._x
Az = 2j * self._k
Axx = 1
Ax = 1 # 1/x-dependency is handled in C++-code
F0 = _compute_potential(n0, self._k)
super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
def step(self, n, boundary):
F = _compute_potential(n, self._k) * self._x
F = _compute_potential(n, self._k)
return super().step(F, boundary)
......
......@@ -72,8 +72,8 @@ namespace finite_differences {
*/
const array_1D step1d_A0F(const Eigen::Ref<const array_1D> rz,
const Eigen::Ref<const array_1D> rxx,
const array_1D step1d_A0F(const scalar rz,
const scalar rxx,
const Eigen::Ref<const array_1D> fp,
const Eigen::Ref<const array_1D> f,
const Eigen::Ref<const array_1D> up,
......@@ -91,17 +91,17 @@ 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(i); // lower
m(1, i-1) = rz(i) + 2. * rxx(i) - f(i); // diagonal
m(2, i-1) = - rxx(i); // upper;
r(i-1) = (rz(i) - 2. * rxx(i) + fp(i)) * up(i)
+ rxx(i) * up(i-1)
+ rxx(i) * up(i+1);
m(0, i-1) = - rxx; // lower
m(1, i-1) = rz + 2. * rxx - f(i); // diagonal
m(2, i-1) = - rxx; // upper;
r(i-1) = (rz - 2. * rxx + fp(i)) * up(i)
+ rxx * up(i-1)
+ rxx * up(i+1);
}
// boundary
r(0) += rxx(1) * u(0);
r(n-1) += rxx(n) * u(n+1);
r(0) += rxx * u(0);
r(n-1) += rxx * u(n+1);
// initialize writeable copy
array_1D u_new {u};
......@@ -110,10 +110,10 @@ namespace finite_differences {
return u_new;
}
const array_1D step1d_AAF(const Eigen::Ref<const array_1D> rz,
const Eigen::Ref<const array_1D> rxx,
const Eigen::Ref<const array_1D> rx,
const array_1D step1d_AAF_sym(const complex rz,
const complex rxx,
const complex rx,
const Eigen::Ref<const array_1D> fp,
const Eigen::Ref<const array_1D> f,
const Eigen::Ref<const array_1D> up,
......@@ -121,125 +121,50 @@ namespace finite_differences {
// _p : i
// _ : i+1
size_t n = u.size()-2;
size_t n = u.size()-1;
// setup tridiagonal n x n matrix M = tridiag(a,b,c) and rhs r
algebra::TriMatrix<scalar> m = algebra::TriMatrix<scalar>::Zero(3, n);
array_1D r = array_1D::Zero(n, 1);
for (size_t i = 1; i <= n; ++i) {
m(0, i-1) = -rxx(i) + rx(i); // lower
m(1, i-1) = rz(i) + 2. * rxx(i) - f(i); // diagonal
m(2, i-1) = -rxx(i) - rx(i); // upper;
r(i-1) = (rz(i) - 2. * rxx(i) + f(i)) * up(i)
+ (rxx(i) - rx(i)) * up(i-1)
+ (rxx(i) + rx(i)) * up(i+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. * rxx - f(i); // diagonal
m(2, i) = -rxx - rhix; // upper;
r(i) = (rz - 2. * rxx + f(i)) * up(i)
+ (rxx - rhix) * up(i-1)
+ (rxx + rhix) * up(i+1);
}
// left boundary
m(0, 0) = 0; // lower
m(1, 0) = rz + 2. * rxx + 2. * rx - f(0); // diagonal
m(2, 0) = -2. * rxx - 2. * rx; // upper;
r(0) = (rz - 2. * rxx - 2. * rx + f(0)) * up(0)
+ (2. * rxx + 2. * rx) * up(1);
// boundary
r(0) += (rxx(1) - rx(1)) * u(0);
r(n-1) += (rxx(n) + rx(n)) * u(n+1);
// right boundary
complex rxnh = rx / static_cast<complex>(n-1) / 2.0;
r(n-1) += (rxx + rxnh) * u(n);
// initialize writeable copy
array_1D u_new {u};
u_new.segment(1, n) = algebra::tridiagonal(m, r);
return u_new;
}
// Two-step alternating-direction Peaceman-Rachford scheme (J. W. Thomas: Numerical Partial Differential Equations)
const array_2D step2d_A0F(
const Eigen::Ref<const array_2D> &rz,
const Eigen::Ref<const array_2D> &rxx,
const Eigen::Ref<const array_2D> &ryy,
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
size_t nx = u.rows() - 2;
size_t ny = u.cols() - 2;
// first halfstep
// x at i + 1/2
// y at i
array_2D uhalf = array_2D::Zero(u.rows(), u.cols());
#pragma omp parallel for
for (size_t iy = 1; iy <= ny; ++iy) {
algebra::TriMatrix<scalar> m = algebra::TriMatrix<scalar>::Zero(3, nx);
array_1D r = array_1D::Zero(nx, 1);
for (size_t ix = 1; ix <= nx; ++ix) {
m(0, ix-1) = - rxx(ix - 1, iy);
m(1, ix-1) = rz(ix, iy) + 2. * rxx(ix, iy) - f(ix, iy);
m(2, ix-1) = - rxx(ix, iy);
r(ix-1) = (rz(ix, iy) - 2. * ryy(ix, iy) + fp(ix, iy)) * up(ix, iy)
+ ryy(ix, iy) * up(ix, iy - 1)
+ ryy(ix, iy) * up(ix, iy + 1);
}
// boundary conditions
// TODO(Leon): use interpolated values?
r(0) += rxx(1, iy) * u(0, iy);
r(nx-1) += rxx(nx, iy) * u(nx+1, iy);
uhalf.col(iy).segment(1, nx) = algebra::tridiagonal(m, r);
}
// second halfstep
// note: internal storage is column-major such that column-wise access would be much faster
array_2D u_new_trans = array_2D::Zero(u.cols(), u.rows()); // transposed shape
#pragma omp parallel for
for (size_t ix = 1; ix <= nx; ++ix) {
algebra::TriMatrix<scalar> m = algebra::TriMatrix<scalar>::Zero(3, ny);
array_1D r = array_1D::Zero(ny, 1);
for (size_t iy = 1; iy <= ny; ++iy) {
m(0, iy-1) = - ryy(ix, iy-1);
m(1, iy-1) = rz(ix, iy) + 2. * ryy(ix, iy) - f(ix, iy);
m(2, iy-1) = - ryy(ix, iy);
r(iy-1) = (rz(ix, iy) - 2. * rxx(ix, iy) + fp(ix, iy)) * uhalf(ix, iy)
+ rxx(ix, iy) * uhalf(ix-1, iy)
+ rxx(ix, iy) * uhalf(ix+1, iy);
}
r(0) += ryy(ix, 1) * u(ix, 0);
r(ny-1) += ryy(ix, ny) * u(ix, ny+1);
u_new_trans.col(ix).segment(1, ny) = algebra::tridiagonal(m, r);
}
array_2D u_new {u_new_trans.transpose()};
u_new.segment(0, n) = algebra::tridiagonal(m, r);
return u_new;
}
// Scalar Version
/* Two-step alternating-direction Peaceman-Rachford scheme
(J. W. Thomas: Numerical Partial Differential Equations)
*/
const array_2D step2d_A0Fs(
const array_2D step2d_A0F(
const complex rz,
const complex rxx,
const complex ryy,
......
......@@ -28,18 +28,8 @@ PYBIND11_MODULE(_solver, m) {
py::arg("up").noconvert(),
py::arg("u").noconvert() );
m.def("step2d_A0Fs",
&fd::step2d_A0Fs,
py::arg("rz").noconvert(),
py::arg("rxx").noconvert(),
py::arg("ryy").noconvert(),
py::arg("fp").noconvert(),
py::arg("f").noconvert(),
py::arg("up").noconvert(),
py::arg("u").noconvert() );
m.def("step1d_AAF",
&fd::step1d_AAF,
m.def("step1d_AAF_sym",
&fd::step1d_AAF_sym,
py::arg("rz").noconvert(),
py::arg("rxx").noconvert(),
py::arg("rx").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