Skip to content

Commit 5996cd3

Browse files
WeiqunZhangWeiqunZhang
committed
WarpX Poisson solver: Option to use staircases approximation
Add an option to MLEBNodeFDLaplacian (i.e., WarpX Poisson solver) to allow for using staircases approximation for EB.
1 parent ebd9a11 commit 5996cd3

5 files changed

+689
-61
lines changed

‎Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLap_2D_K.H

Lines changed: 266 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,73 @@ void mlebndfdlap_adotx_eb (int i, int j, int k, Array4<Real> const& y,
9595
bx, by);
9696
}
9797

98+
template <typename F>
99+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
100+
void mlebndfdlap_adotx_eb_sc_doit (int i, int j, int k, Array4<Real> const& y,
101+
Array4<Real const> const& x, Array4<Real const> const& levset,
102+
Array4<int const> const& dmsk,
103+
F const& xeb, Real bx, Real by) noexcept
104+
{
105+
if (dmsk(i,j,k)) {
106+
y(i,j,k) = Real(0.0);
107+
} else {
108+
Real tmp, out;
109+
110+
if (levset(i+1,j,k) < Real(0.0)) {
111+
tmp = x(i+1,j,k) - x(i,j,k);
112+
} else {
113+
tmp = (xeb(i+1,j,k) - x(i,j,k));
114+
}
115+
116+
if (levset(i-1,j,k) < Real(0.0)) {
117+
tmp += x(i-1,j,k) - x(i,j,k);
118+
} else {
119+
tmp += (xeb(i-1,j,k) - x(i,j,k));
120+
}
121+
122+
out = tmp * bx;
123+
124+
if (levset(i,j+1,k) < Real(0.0)) {
125+
tmp = x(i,j+1,k) - x(i,j,k);
126+
} else {
127+
tmp = (xeb(i,j+1,k) - x(i,j,k));
128+
}
129+
130+
if (levset(i,j-1,k) < Real(0.0)) {
131+
tmp += x(i,j-1,k) - x(i,j,k);
132+
} else {
133+
tmp += (xeb(i,j-1,k) - x(i,j,k));
134+
}
135+
136+
out += tmp * by;
137+
138+
y(i,j,k) = out;
139+
}
140+
}
141+
142+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
143+
void mlebndfdlap_adotx_eb_sc (int i, int j, int k, Array4<Real> const& y,
144+
Array4<Real const> const& x, Array4<Real const> const& levset,
145+
Array4<int const> const& dmsk,
146+
Real xeb, Real bx, Real by) noexcept
147+
{
148+
mlebndfdlap_adotx_eb_sc_doit(i, j, k, y, x, levset, dmsk,
149+
[=] (int, int, int) -> Real { return xeb; },
150+
bx, by);
151+
}
152+
153+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
154+
void mlebndfdlap_adotx_eb_sc (int i, int j, int k, Array4<Real> const& y,
155+
Array4<Real const> const& x, Array4<Real const> const& levset,
156+
Array4<int const> const& dmsk,
157+
Array4<Real const> const& xeb, Real bx, Real by) noexcept
158+
{
159+
mlebndfdlap_adotx_eb_sc_doit(i, j, k, y, x, levset, dmsk,
160+
[=] (int i1, int i2, int i3) -> Real {
161+
return xeb(i1,i2,i3); },
162+
bx, by);
163+
}
164+
98165
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
99166
void mlebndfdlap_adotx (int i, int j, int k, Array4<Real> const& y,
100167
Array4<Real const> const& x, Array4<int const> const& dmsk,
@@ -172,6 +239,51 @@ void mlebndfdlap_gsrb_eb (int i, int j, int k, Array4<Real> const& x,
172239
}
173240
}
174241

242+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
243+
void mlebndfdlap_gsrb_eb_sc (int i, int j, int k, Array4<Real> const& x,
244+
Array4<Real const> const& rhs, Array4<Real const> const& levset,
245+
Array4<int const> const& dmsk,
246+
Real bx, Real by, int redblack) noexcept
247+
{
248+
if ((i+j+k+redblack)%2 == 0) {
249+
if (dmsk(i,j,k)) {
250+
x(i,j,k) = Real(0.);
251+
} else {
252+
Real tmp1;
253+
254+
if (levset(i+1,j,k) < Real(0.0)) { // regular
255+
tmp1 = x(i+1,j,k);
256+
} else {
257+
tmp1 = Real(0.0);
258+
}
259+
260+
if (levset(i-1,j,k) < Real(0.0)) {
261+
tmp1 += x(i-1,j,k);
262+
}
263+
264+
Real gamma = Real(-2.0) * bx;
265+
Real rho = tmp1 * bx;
266+
267+
if (levset(i,j+1,k) < Real(0.0)) {
268+
tmp1 = x(i,j+1,k);
269+
} else {
270+
tmp1 = Real(0.0);
271+
}
272+
273+
if (levset(i,j-1,k) < Real(0.0)) {
274+
tmp1 += x(i,j-1,k);
275+
}
276+
277+
gamma += Real(-2.0) * by;
278+
rho += tmp1 * by;
279+
280+
Real Ax = rho + gamma*x(i,j,k);
281+
constexpr Real omega = Real(1.25);
282+
x(i,j,k) += (rhs(i,j,k) - Ax) * (omega / (gamma));
283+
}
284+
}
285+
}
286+
175287
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
176288
void mlebndfdlap_gsrb (int i, int j, int k, Array4<Real> const& x,
177289
Array4<Real const> const& rhs, Array4<int const> const& dmsk,
@@ -288,6 +400,88 @@ void mlebndfdlap_adotx_rz_eb (int i, int j, int k, Array4<Real> const& y,
288400
sigr, dr, dz, rlo, alpha);
289401
}
290402

403+
template <typename F>
404+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
405+
void mlebndfdlap_adotx_rz_eb_sc_doit (int i, int j, int k, Array4<Real> const& y,
406+
Array4<Real const> const& x, Array4<Real const> const& levset,
407+
Array4<int const> const& dmsk,
408+
F const& xeb, Real sigr, Real dr, Real dz, Real rlo, Real alpha) noexcept
409+
{
410+
Real const r = rlo + Real(i) * dr;
411+
if (dmsk(i,j,k) || (r == Real(0.0) && alpha != Real(0.0))) {
412+
y(i,j,k) = Real(0.0);
413+
} else {
414+
Real tmp, out;
415+
416+
if (r == Real(0.0)) {
417+
if (levset(i+1,j,k) < Real(0.0)) { // regular
418+
out = Real(4.0) * sigr * (x(i+1,j,k)-x(i,j,k)) / (dr*dr);
419+
} else {
420+
out = Real(4.0) * sigr * (xeb(i+1,j,k)-x(i,j,k)) / (dr*dr);
421+
}
422+
} else {
423+
if (levset(i+1,j,k) < Real(0.0)) { // regular
424+
tmp = (x(i+1,j,k) - x(i,j,k)) * (r + Real(0.5) * dr);
425+
} else {
426+
tmp = (xeb(i+1,j,k) - x(i,j,k)) * (r + Real(0.5) * dr);
427+
}
428+
429+
if (levset(i-1,j,k) < Real(0.0)) {
430+
tmp += (x(i-1,j,k) - x(i,j,k)) * (r - Real(0.5) * dr);
431+
} else {
432+
tmp += (xeb(i-1,j,k) - x(i,j,k)) * (r - Real(0.5) * dr);
433+
}
434+
435+
out = tmp * sigr / (r * dr * dr);
436+
}
437+
438+
if (levset(i,j+1,k) < Real(0.0)) {
439+
tmp = x(i,j+1,k) - x(i,j,k);
440+
} else {
441+
tmp = (xeb(i,j+1,k) - x(i,j,k));
442+
}
443+
444+
445+
if (levset(i,j-1,k) < Real(0.0)) {
446+
tmp += x(i,j-1,k) - x(i,j,k);
447+
} else {
448+
tmp += (xeb(i,j-1,k) - x(i,j,k));
449+
}
450+
451+
out += tmp / (dz * dz);
452+
453+
if (r != Real(0.0)) {
454+
out -= alpha/(r*r) * x(i,j,k);
455+
}
456+
457+
y(i,j,k) = out;
458+
}
459+
}
460+
461+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
462+
void mlebndfdlap_adotx_rz_eb_sc (int i, int j, int k, Array4<Real> const& y,
463+
Array4<Real const> const& x, Array4<Real const> const& levset,
464+
Array4<int const> const& dmsk,
465+
Real xeb, Real sigr, Real dr, Real dz, Real rlo, Real alpha) noexcept
466+
{
467+
mlebndfdlap_adotx_rz_eb_sc_doit(i, j, k, y, x, levset, dmsk,
468+
[=] (int, int, int) -> Real { return xeb; },
469+
sigr, dr, dz, rlo, alpha);
470+
}
471+
472+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
473+
void mlebndfdlap_adotx_rz_eb_sc (int i, int j, int k, Array4<Real> const& y,
474+
Array4<Real const> const& x, Array4<Real const> const& levset,
475+
Array4<int const> const& dmsk,
476+
Array4<Real const> const& xeb, Real sigr, Real dr, Real dz, Real rlo,
477+
Real alpha) noexcept
478+
{
479+
mlebndfdlap_adotx_rz_eb_sc_doit(i, j, k, y, x, levset, dmsk,
480+
[=] (int i1, int i2, int i3) -> Real {
481+
return xeb(i1,i2,i3); },
482+
sigr, dr, dz, rlo, alpha);
483+
}
484+
291485
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
292486
void mlebndfdlap_adotx_rz (int i, int j, int k, Array4<Real> const& y,
293487
Array4<Real const> const& x, Array4<int const> const& dmsk,
@@ -393,6 +587,78 @@ void mlebndfdlap_gsrb_rz_eb (int i, int j, int k, Array4<Real> const& x,
393587
}
394588
}
395589

590+
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
591+
void mlebndfdlap_gsrb_rz_eb_sc (int i, int j, int k, Array4<Real> const& x,
592+
Array4<Real const> const& rhs, Array4<Real const> const& levset,
593+
Array4<int const> const& dmsk,
594+
Real sigr, Real dr, Real dz, Real rlo, int redblack, Real alpha) noexcept
595+
{
596+
if ((i+j+k+redblack)%2 == 0) {
597+
Real const r = rlo + Real(i) * dr;
598+
if (dmsk(i,j,k) || (r == Real(0.0) && alpha != Real(0.0))) {
599+
x(i,j,k) = Real(0.);
600+
} else {
601+
Real tmp, tmp0, Ax, gamma;
602+
603+
if (r == Real(0.0)) {
604+
if (levset(i+1,j,k) < Real(0.0)) { // regular
605+
Ax = (Real(4.0) * sigr / (dr*dr)) * (x(i+1,j,k)-x(i,j,k));
606+
gamma = -(Real(4.0) * sigr / (dr*dr));
607+
} else {
608+
gamma = -(Real(4.0) * sigr / (dr*dr));
609+
Ax = gamma * x(i,j,k);
610+
}
611+
} else {
612+
if (levset(i+1,j,k) < Real(0.0)) { // regular
613+
tmp = (x(i+1,j,k) - x(i,j,k)) * (r + Real(0.5) * dr);
614+
tmp0 = -(r + Real(0.5) * dr);
615+
} else {
616+
tmp0 = Real(-1.0) * (r + Real(0.5) * dr);
617+
tmp = tmp0 * x(i,j,k);
618+
}
619+
620+
if (levset(i-1,j,k) < Real(0.0)) {
621+
tmp += (x(i-1,j,k) - x(i,j,k)) * (r - Real(0.5) * dr);
622+
tmp0 += -(r - Real(0.5) * dr);
623+
} else {
624+
tmp += -x(i,j,k) * (r - Real(0.5) * dr);
625+
tmp0 += Real(-1.0) * (r - Real(0.5) * dr);
626+
}
627+
628+
Ax = tmp * sigr / (r * dr * dr);
629+
gamma = tmp0 * sigr / (r * dr * dr);
630+
}
631+
632+
if (levset(i,j+1,k) < Real(0.0)) {
633+
tmp = x(i,j+1,k) - x(i,j,k);
634+
tmp0 = Real(-1.0);
635+
} else {
636+
tmp0 = Real(-1.0);
637+
tmp = tmp0 * x(i,j,k);
638+
}
639+
640+
if (levset(i,j-1,k) < Real(0.0)) {
641+
tmp += x(i,j-1,k) - x(i,j,k);
642+
tmp0 += Real(-1.0);
643+
} else {
644+
tmp += -x(i,j,k);
645+
tmp0 += Real(-1.0);
646+
}
647+
648+
Ax += tmp / (dz * dz);
649+
gamma += tmp0 / (dz * dz);
650+
651+
if (r != Real(0.0)) {
652+
Ax -= alpha/(r*r) * x(i,j,k);
653+
gamma -= alpha/(r*r);
654+
}
655+
656+
constexpr Real omega = Real(1.25);
657+
x(i,j,k) += (rhs(i,j,k) - Ax) * (omega / (gamma));
658+
}
659+
}
660+
}
661+
396662
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
397663
void mlebndfdlap_gsrb_rz (int i, int j, int k, Array4<Real> const& x,
398664
Array4<Real const> const& rhs, Array4<int const> const& dmsk,

0 commit comments

Comments
 (0)