Skip to content

Commit 4fb7103

Browse files
WeiqunZhangWeiqunZhang
authored
FFT Poisson Solver: Support domain not starting with 0 (#4573)
1 parent 850a0fb commit 4fb7103

File tree

1 file changed

+84
-23
lines changed

1 file changed

+84
-23
lines changed

‎Src/FFT/AMReX_FFT_Poisson.H‎

Lines changed: 84 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,42 @@ namespace detail {
1111
template <typename MF>
1212
void fill_physbc (MF& mf, Geometry const& geom,
1313
Array<std::pair<Boundary,Boundary>,AMREX_SPACEDIM> const& bc);
14+
15+
inline Geometry shift_geom (Geometry const& geom)
16+
{
17+
auto const& dom = geom.Domain();
18+
auto const& dlo = dom.smallEnd();
19+
if (dlo == 0) {
20+
return geom;
21+
} else {
22+
return Geometry(amrex::shift(dom,-dlo), geom.ProbDomain(), geom.CoordInt(),
23+
geom.isPeriodic());
24+
}
25+
}
26+
27+
template <typename MF>
28+
void shift_mfs (IntVect const& domain_lo, MF const& mf1, MF const& mf2,
29+
MF& new_mf1, MF& new_mf2)
30+
{
31+
AMREX_ALWAYS_ASSERT(mf1.boxArray() == mf2.boxArray() &&
32+
mf1.DistributionMap() == mf2.DistributionMap());
33+
BoxArray ba = mf1.boxArray();
34+
ba.shift(-domain_lo);
35+
new_mf1.define(ba, mf1.DistributionMap(), 1, mf1.nGrowVect(),
36+
MFInfo().SetAlloc(false));
37+
new_mf2.define(ba, mf2.DistributionMap(), 1, mf2.nGrowVect(),
38+
MFInfo().SetAlloc(false));
39+
for (MFIter mfi(mf1); mfi.isValid(); ++mfi) {
40+
typename MF::fab_type fab1(mf1[mfi], amrex::make_alias, 0, 1);
41+
fab1.shift(-domain_lo);
42+
new_mf1.setFab(mfi, std::move(fab1));
43+
44+
typename MF::fab_type fab2(mf2[mfi], amrex::make_alias, 0, 1);
45+
fab2.shift(-domain_lo);
46+
new_mf2.setFab(mfi, std::move(fab2));
47+
}
48+
}
49+
1450
}
1551

1652
/**
@@ -25,7 +61,9 @@ public:
2561
template <typename FA=MF, std::enable_if_t<IsFabArray_v<FA>,int> = 0>
2662
Poisson (Geometry const& geom,
2763
Array<std::pair<Boundary,Boundary>,AMREX_SPACEDIM> const& bc)
28-
: m_geom(geom), m_bc(bc)
64+
: m_domain_lo(geom.Domain().smallEnd()),
65+
m_geom(detail::shift_geom(geom)),
66+
m_bc(bc)
2967
{
3068
bool all_periodic = true;
3169
for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
@@ -42,7 +80,8 @@ public:
4280

4381
template <typename FA=MF, std::enable_if_t<IsFabArray_v<FA>,int> = 0>
4482
explicit Poisson (Geometry const& geom)
45-
: m_geom(geom),
83+
: m_domain_lo(geom.Domain().smallEnd()),
84+
m_geom(detail::shift_geom(geom)),
4685
m_bc{AMREX_D_DECL(std::make_pair(Boundary::periodic,Boundary::periodic),
4786
std::make_pair(Boundary::periodic,Boundary::periodic),
4887
std::make_pair(Boundary::periodic,Boundary::periodic))}
@@ -61,9 +100,10 @@ public:
61100
* except for the corners of the physical domain where they are not used
62101
* in the cross stencil of the operator. The two MFs could be the same MF.
63102
*/
64-
void solve (MF& soln, MF const& rhs);
103+
void solve (MF& a_soln, MF const& a_rhs);
65104

66105
private:
106+
IntVect m_domain_lo;
67107
Geometry m_geom;
68108
Array<std::pair<Boundary,Boundary>,AMREX_SPACEDIM> m_bc;
69109
std::unique_ptr<R2X<typename MF::value_type>> m_r2x;
@@ -110,7 +150,9 @@ public:
110150
template <typename FA=MF, std::enable_if_t<IsFabArray_v<FA>,int> = 0>
111151
PoissonHybrid (Geometry const& geom,
112152
Array<std::pair<Boundary,Boundary>,AMREX_SPACEDIM> const& bc)
113-
: m_geom(geom), m_bc(bc)
153+
: m_domain_lo(geom.Domain().smallEnd()),
154+
m_geom(detail::shift_geom(geom)),
155+
m_bc(bc)
114156
{
115157
#if (AMREX_SPACEDIM < 3)
116158
amrex::Abort("FFT::PoissonHybrid: 1D & 2D todo");
@@ -152,7 +194,7 @@ public:
152194
void solve (MF& soln, MF const& rhs, Gpu::DeviceVector<T> const& dz);
153195

154196
template <typename TRIA, typename TRIC>
155-
void solve (MF& soln, MF const& rhs, TRIA const& tria, TRIC const& tric);
197+
void solve (MF& a_soln, MF const& a_rhs, TRIA const& tria, TRIC const& tric);
156198

157199
// This is public for cuda
158200
template <typename FA, typename TRIA, typename TRIC>
@@ -164,6 +206,7 @@ private:
164206

165207
void build_spmf ();
166208

209+
IntVect m_domain_lo;
167210
Geometry m_geom;
168211
Array<std::pair<Boundary,Boundary>,AMREX_SPACEDIM> m_bc;
169212
std::unique_ptr<R2X<typename MF::value_type>> m_r2x;
@@ -174,11 +217,20 @@ private:
174217
};
175218

176219
template <typename MF>
177-
void Poisson<MF>::solve (MF& soln, MF const& rhs)
220+
void Poisson<MF>::solve (MF& a_soln, MF const& a_rhs)
178221
{
179222
BL_PROFILE("FFT::Poisson::solve");
180223

181-
AMREX_ASSERT(soln.is_cell_centered() && rhs.is_cell_centered());
224+
MF* soln = &a_soln;
225+
MF const* rhs = &a_rhs;
226+
MF solntmp, rhstmp;
227+
if (m_domain_lo != 0) {
228+
detail::shift_mfs(m_domain_lo, a_soln, a_rhs, solntmp, rhstmp);
229+
soln = &solntmp;
230+
rhs = &rhstmp;
231+
}
232+
233+
AMREX_ASSERT(soln->is_cell_centered() && rhs->is_cell_centered());
182234

183235
using T = typename MF::value_type;
184236

@@ -233,15 +285,15 @@ void Poisson<MF>::solve (MF& soln, MF const& rhs)
233285
spectral_data *= scale;
234286
};
235287

236-
IntVect const& ng = amrex::elemwiseMin(soln.nGrowVect(), IntVect(1));
288+
IntVect const& ng = amrex::elemwiseMin(soln->nGrowVect(), IntVect(1));
237289

238290
if (m_r2x) {
239-
m_r2x->forwardThenBackward_doit_0(rhs, soln, f, ng, m_geom.periodicity());
240-
detail::fill_physbc(soln, m_geom, m_bc);
291+
m_r2x->forwardThenBackward_doit_0(*rhs, *soln, f, ng, m_geom.periodicity());
292+
detail::fill_physbc(*soln, m_geom, m_bc);
241293
} else {
242-
m_r2c->forward(rhs);
294+
m_r2c->forward(*rhs);
243295
m_r2c->post_forward_doit_0(f);
244-
m_r2c->backward_doit(soln, ng, m_geom.periodicity());
296+
m_r2c->backward_doit(*soln, ng, m_geom.periodicity());
245297
}
246298
}
247299

@@ -430,39 +482,48 @@ void PoissonHybrid<MF>::solve (MF& soln, MF const& rhs, Vector<T> const& dz)
430482

431483
template <typename MF>
432484
template <typename TRIA, typename TRIC>
433-
void PoissonHybrid<MF>::solve (MF& soln, MF const& rhs, TRIA const& tria,
485+
void PoissonHybrid<MF>::solve (MF& a_soln, MF const& a_rhs, TRIA const& tria,
434486
TRIC const& tric)
435487
{
436488
BL_PROFILE("FFT::PoissonHybrid::solve");
437489

438-
AMREX_ASSERT(soln.is_cell_centered() && rhs.is_cell_centered());
490+
AMREX_ASSERT(a_soln.is_cell_centered() && a_rhs.is_cell_centered());
439491

440492
#if (AMREX_SPACEDIM < 3)
441-
amrex::ignore_unused(soln, rhs, tria, tric);
493+
amrex::ignore_unused(a_soln, a_rhs, tria, tric);
442494
#else
443495

444-
IntVect const& ng = amrex::elemwiseMin(soln.nGrowVect(), IntVect(1));
496+
MF* soln = &a_soln;
497+
MF const* rhs = &a_rhs;
498+
MF solntmp, rhstmp;
499+
if (m_domain_lo != 0) {
500+
detail::shift_mfs(m_domain_lo, a_soln, a_rhs, solntmp, rhstmp);
501+
soln = &solntmp;
502+
rhs = &rhstmp;
503+
}
504+
505+
IntVect const& ng = amrex::elemwiseMin(soln->nGrowVect(), IntVect(1));
445506

446507
if (m_r2c)
447508
{
448-
m_r2c->forward(rhs, m_spmf_c);
509+
m_r2c->forward(*rhs, m_spmf_c);
449510
solve_z(m_spmf_c, tria, tric);
450-
m_r2c->backward_doit(m_spmf_c, soln, ng, m_geom.periodicity());
511+
m_r2c->backward_doit(m_spmf_c, *soln, ng, m_geom.periodicity());
451512
}
452513
else
453514
{
454515
if (m_r2x->m_cy.empty()) { // spectral data is real
455-
m_r2x->forward(rhs, m_spmf_r);
516+
m_r2x->forward(*rhs, m_spmf_r);
456517
solve_z(m_spmf_r, tria, tric);
457-
m_r2x->backward(m_spmf_r, soln, ng, m_geom.periodicity());
518+
m_r2x->backward(m_spmf_r, *soln, ng, m_geom.periodicity());
458519
} else { // spectral data is complex.
459-
m_r2x->forward(rhs, m_spmf_c);
520+
m_r2x->forward(*rhs, m_spmf_c);
460521
solve_z(m_spmf_c, tria, tric);
461-
m_r2x->backward(m_spmf_c, soln, ng, m_geom.periodicity());
522+
m_r2x->backward(m_spmf_c, *soln, ng, m_geom.periodicity());
462523
}
463524
}
464525

465-
detail::fill_physbc(soln, m_geom, m_bc);
526+
detail::fill_physbc(*soln, m_geom, m_bc);
466527
#endif
467528
}
468529

0 commit comments

Comments
 (0)