Skip to content

Commit

Permalink
2D RZ solver for WarpX: Arbitrary coefficient (AMReX-Codes#2986)
Browse files Browse the repository at this point in the history 
The assumption in the 2D RZ solver for WarpX used to be there was no
sigma_r (i.e., sigma_r == 1). In this PR, we allow arbitrary sigma_r
coefficient.
  • Loading branch information
@WeiqunZhang
WeiqunZhang authored Oct 15, 2022
1 parent 9a3cd5d commit bcbf17f
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 31 deletions.
44 changes: 22 additions & 22 deletions Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLap_2D_K.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void mlebndfdlap_adotx_rz_eb_doit (int i, int j, int k, Array4<Real> const& y,
Array4<Real const> const& x, Array4<int const> const& dmsk,
Array4<Real const> const& ecx, Array4<Real const> const& ecy,
F && xeb, Real dr, Real dz, Real rlo) noexcept
F && xeb, Real sigr, Real dr, Real dz, Real rlo) noexcept
{
if (dmsk(i,j,k)) {
y(i,j,k) = Real(0.0);
Expand All @@ -211,11 +211,11 @@ void mlebndfdlap_adotx_rz_eb_doit (int i, int j, int k, Array4<Real> const& y,
Real const r = rlo + Real(i) * dr;
if (r == Real(0.0)) {
if (ecx(i,j,k) == Real(1.0)) { // regular
out = Real(4.0) * (x(i+1,j,k)-x(i,j,k)) / (dr*dr);
out = Real(4.0) * sigr * (x(i+1,j,k)-x(i,j,k)) / (dr*dr);
scale = Real(1.0);
} else {
hp = Real(1.0) + Real(2.) * ecx(i,j,k);
out = Real(4.0) * (xeb(i+1,j,k)-x(i,j,k)) / (dr*dr*hp*hp);
out = Real(4.0) * sigr * (xeb(i+1,j,k)-x(i,j,k)) / (dr*dr*hp*hp);
scale = hp;
}
} else {
Expand All @@ -235,7 +235,7 @@ void mlebndfdlap_adotx_rz_eb_doit (int i, int j, int k, Array4<Real> const& y,
tmp += (xeb(i-1,j,k) - x(i,j,k)) / hm * (r - Real(0.5) * hp * dr);
}

out = tmp * Real(2.0) / ((hp+hm) * r * dr * dr);
out = tmp * Real(2.0) * sigr / ((hp+hm) * r * dr * dr);
scale = amrex::min(hm, hp);
}

Expand Down Expand Up @@ -266,41 +266,41 @@ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void mlebndfdlap_adotx_rz_eb (int i, int j, int k, Array4<Real> const& y,
Array4<Real const> const& x, Array4<int const> const& dmsk,
Array4<Real const> const& ecx, Array4<Real const> const& ecy,
Real xeb, Real dr, Real dz, Real rlo) noexcept
Real xeb, Real sigr, Real dr, Real dz, Real rlo) noexcept
{
mlebndfdlap_adotx_rz_eb_doit(i, j, k, y, x, dmsk, ecx, ecy,
[=] (int, int, int) -> Real { return xeb; },
dr, dz, rlo);
sigr, dr, dz, rlo);
}

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void mlebndfdlap_adotx_rz_eb (int i, int j, int k, Array4<Real> const& y,
Array4<Real const> const& x, Array4<int const> const& dmsk,
Array4<Real const> const& ecx, Array4<Real const> const& ecy,
Array4<Real const> const& xeb, Real dr, Real dz, Real rlo) noexcept
Array4<Real const> const& xeb, Real sigr, Real dr, Real dz, Real rlo) noexcept
{
mlebndfdlap_adotx_rz_eb_doit(i, j, k, y, x, dmsk, ecx, ecy,
[=] (int i1, int i2, int i3) -> Real {
return xeb(i1,i2,i3); },
dr, dz, rlo);
sigr, dr, dz, rlo);
}

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void mlebndfdlap_adotx_rz (int i, int j, int k, Array4<Real> const& y,
Array4<Real const> const& x, Array4<int const> const& dmsk,
Real dr, Real dz, Real rlo) noexcept
Real sigr, Real dr, Real dz, Real rlo) noexcept
{
if (dmsk(i,j,k)) {
y(i,j,k) = Real(0.0);
} else {
Real Ax = (x(i,j-1,k) - Real(2.0)*x(i,j,k) + x(i,j+1,k)) / (dz*dz);
Real const r = rlo + Real(i)*dr;
if (r == Real(0.0)) {
Ax += Real(4.0) * (x(i+1,j,k)-x(i,j,k)) / (dr*dr);
Ax += Real(4.0) * sigr * (x(i+1,j,k)-x(i,j,k)) / (dr*dr);
} else {
Real const rp = r + Real(0.5)*dr;
Real const rm = r - Real(0.5)*dr;
Ax += (rp*x(i+1,j,k) - (rp+rm)*x(i,j,k) + rm*x(i-1,j,k)) / (r*dr*dr);
Ax += sigr * (rp*x(i+1,j,k) - (rp+rm)*x(i,j,k) + rm*x(i-1,j,k)) / (r*dr*dr);
}
y(i,j,k) = Ax;
}
Expand All @@ -310,7 +310,7 @@ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void mlebndfdlap_gsrb_rz_eb (int i, int j, int k, Array4<Real> const& x,
Array4<Real const> const& rhs, Array4<int const> const& dmsk,
Array4<Real const> const& ecx, Array4<Real const> const& ecy,
Real dr, Real dz, Real rlo, int redblack) noexcept
Real sigr, Real dr, Real dz, Real rlo, int redblack) noexcept
{
if ((i+j+k+redblack)%2 == 0) {
if (dmsk(i,j,k)) {
Expand All @@ -322,12 +322,12 @@ void mlebndfdlap_gsrb_rz_eb (int i, int j, int k, Array4<Real> const& x,
Real const r = rlo + Real(i) * dr;
if (r == Real(0.0)) {
if (ecx(i,j,k) == Real(1.0)) { // regular
Ax = (Real(4.0) / (dr*dr)) * (x(i+1,j,k)-x(i,j,k));
gamma = -(Real(4.0) / (dr*dr));
Ax = (Real(4.0) * sigr / (dr*dr)) * (x(i+1,j,k)-x(i,j,k));
gamma = -(Real(4.0) * sigr / (dr*dr));
scale = Real(1.0);
} else {
hp = Real(1.0) + Real(2.) * ecx(i,j,k);
gamma = -(Real(4.0) / (dr*dr*hp*hp));
gamma = -(Real(4.0) * sigr / (dr*dr*hp*hp));
Ax = gamma * x(i,j,k);
scale = hp;
}
Expand All @@ -352,8 +352,8 @@ void mlebndfdlap_gsrb_rz_eb (int i, int j, int k, Array4<Real> const& x,
tmp0 += Real(-1.0) / hm * (r - Real(0.5) * hp * dr);
}

Ax = tmp * Real(2.0) / ((hp+hm) * r * dr * dr);
gamma = tmp0 * Real(2.0) / ((hp+hm) * r * dr * dr);
Ax = tmp * Real(2.0) * sigr / ((hp+hm) * r * dr * dr);
gamma = tmp0 * Real(2.0) * sigr / ((hp+hm) * r * dr * dr);
scale = amrex::min(hm, hp);
}

Expand Down Expand Up @@ -390,7 +390,7 @@ void mlebndfdlap_gsrb_rz_eb (int i, int j, int k, Array4<Real> const& x,
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void mlebndfdlap_gsrb_rz (int i, int j, int k, Array4<Real> const& x,
Array4<Real const> const& rhs, Array4<int const> const& dmsk,
Real dr, Real dz, Real rlo, int redblack) noexcept
Real sigr, Real dr, Real dz, Real rlo, int redblack) noexcept
{
if ((i+j+k+redblack)%2 == 0) {
if (dmsk(i,j,k)) {
Expand All @@ -400,13 +400,13 @@ void mlebndfdlap_gsrb_rz (int i, int j, int k, Array4<Real> const& x,
Real gamma = -Real(2.0) / (dz*dz);
Real const r = rlo + Real(i)*dr;
if (r == Real(0.0)) {
Ax += (Real(4.0)/(dr*dr)) * (x(i+1,j,k)-x(i,j,k));
gamma += -(Real(4.0)/(dr*dr));
Ax += (Real(4.0)*sigr/(dr*dr)) * (x(i+1,j,k)-x(i,j,k));
gamma += -(Real(4.0)*sigr/(dr*dr));
} else {
Real const rp = r + Real(0.5)*dr;
Real const rm = r - Real(0.5)*dr;
Ax += (rp*x(i+1,j,k) - (rp+rm)*x(i,j,k) + rm*x(i-1,j,k)) / (r*dr*dr);
gamma += -(rp+rm) / (r*dr*dr);
Ax += sigr*(rp*x(i+1,j,k) - (rp+rm)*x(i,j,k) + rm*x(i-1,j,k)) / (r*dr*dr);
gamma += -sigr*(rp+rm) / (r*dr*dr);
}
constexpr Real omega = Real(1.25);
x(i,j,k) += (rhs(i,j,k) - Ax) * (omega / gamma);
Expand Down
4 changes: 2 additions & 2 deletions Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLaplacian.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ namespace amrex {
// with only diagonal components. The EB is assumed to be Dirichlet.
//
// del dot (simga grad phi) - alpha/r^2 phi = rhs, for RZ where alpha is a
// scalar constant that is zero by default. sigma is non-zero in
// z-direction only. For now the `alpha` term has not been implemented yet.
// scalar constant that is zero by default. For now the `alpha` term has
// not been implemented yet

class MLEBNodeFDLaplacian
: public MLNodeLinOp
Expand Down
17 changes: 10 additions & 7 deletions Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLaplacian.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -310,8 +310,9 @@ MLEBNodeFDLaplacian::prepareForSolve ()
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_lobc[0][0] == BCType::Neumann,
"The lo-x BC must be Neumann for 2d RZ");
}
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_sigma[0] == 0._rt,
"r-direction sigma must be zero");
if (m_sigma[0] == 0._rt) {
m_sigma[0] = 1._rt; // For backward compatibility
}
}
#endif
}
Expand Down Expand Up @@ -356,6 +357,7 @@ MLEBNodeFDLaplacian::Fapply (int amrlev, int mglev, MultiFab& out, const MultiFa

const auto dxinv = m_geom[amrlev][mglev].InvCellSizeArray();
#if (AMREX_SPACEDIM == 2)
const auto sig0 = m_sigma[0];
const auto dx0 = m_geom[amrlev][mglev].CellSize(0);
const auto dx1 = m_geom[amrlev][mglev].CellSize(1)/std::sqrt(m_sigma[1]);
const auto xlo = m_geom[amrlev][mglev].ProbLo(0);
Expand Down Expand Up @@ -396,7 +398,7 @@ MLEBNodeFDLaplacian::Fapply (int amrlev, int mglev, MultiFab& out, const MultiFa
AMREX_HOST_DEVICE_FOR_3D(box, i, j, k,
{
mlebndfdlap_adotx_rz_eb(i,j,k,yarr,xarr,dmarr,ecx,ecy,
phiebarr, dx0, dx1, xlo);
phiebarr, sig0, dx0, dx1, xlo);
});
} else
#endif
Expand All @@ -413,7 +415,7 @@ MLEBNodeFDLaplacian::Fapply (int amrlev, int mglev, MultiFab& out, const MultiFa
AMREX_HOST_DEVICE_FOR_3D(box, i, j, k,
{
mlebndfdlap_adotx_rz_eb(i,j,k,yarr,xarr,dmarr,ecx,ecy,
phieb, dx0, dx1, xlo);
phieb, sig0, dx0, dx1, xlo);
});
} else
#endif
Expand All @@ -432,7 +434,7 @@ MLEBNodeFDLaplacian::Fapply (int amrlev, int mglev, MultiFab& out, const MultiFa
if (m_rz) {
AMREX_HOST_DEVICE_FOR_3D(box, i, j, k,
{
mlebndfdlap_adotx_rz(i,j,k,yarr,xarr,dmarr,dx0,dx1,xlo);
mlebndfdlap_adotx_rz(i,j,k,yarr,xarr,dmarr,sig0,dx0,dx1,xlo);
});
} else
#endif
Expand All @@ -453,6 +455,7 @@ MLEBNodeFDLaplacian::Fsmooth (int amrlev, int mglev, MultiFab& sol, const MultiF

const auto dxinv = m_geom[amrlev][mglev].InvCellSizeArray();
#if (AMREX_SPACEDIM == 2)
const auto sig0 = m_sigma[0];
const auto dx0 = m_geom[amrlev][mglev].CellSize(0);
const auto dx1 = m_geom[amrlev][mglev].CellSize(1)/std::sqrt(m_sigma[1]);
const auto xlo = m_geom[amrlev][mglev].ProbLo(0);
Expand Down Expand Up @@ -495,7 +498,7 @@ MLEBNodeFDLaplacian::Fsmooth (int amrlev, int mglev, MultiFab& sol, const MultiF
AMREX_HOST_DEVICE_FOR_3D(box, i, j, k,
{
mlebndfdlap_gsrb_rz_eb(i,j,k,solarr,rhsarr,dmskarr,ecx,ecy,
dx0, dx1, xlo, redblack);
sig0, dx0, dx1, xlo, redblack);
});
} else
#endif
Expand All @@ -514,7 +517,7 @@ MLEBNodeFDLaplacian::Fsmooth (int amrlev, int mglev, MultiFab& sol, const MultiF
AMREX_HOST_DEVICE_FOR_3D(box, i, j, k,
{
mlebndfdlap_gsrb_rz(i,j,k,solarr,rhsarr,dmskarr,
dx0, dx1, xlo, redblack);
sig0, dx0, dx1, xlo, redblack);
});
} else
#endif
Expand Down

0 comments on commit bcbf17f

Please sign in to comment.