Skip to content

Commit

Permalink
Add FFT Based MAC Projector
Browse files Browse the repository at this point in the history
  • Loading branch information
WeiqunZhang committed Nov 23, 2024
1 parent 3788361 commit f89137f
Show file tree
Hide file tree
Showing 9 changed files with 274 additions and 62 deletions.
6 changes: 6 additions & 0 deletions Projections/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -12,3 +12,9 @@ target_sources(
hydro_NodalProjector.cpp
hydro_NodalProjector.H
)

if (HYDRO_FFT)
target_sources(amrex_hydro PRIVATE
hydro_FFTMacProjector.cpp
hydro_FFTMacProjector.H)
endif()
5 changes: 5 additions & 0 deletions Projections/Make.package
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,8 @@ CEXE_headers += hydro_NodalProjector.H

CEXE_sources += hydro_MacProjector.cpp
CEXE_sources += hydro_NodalProjector.cpp

ifeq ($(USE_FFT),TRUE)
CEXE_sources += hydro_FFTMacProjector.cpp
CEXE_headers += hydro_FFTMacProjector.H
endif
30 changes: 30 additions & 0 deletions Projections/hydro_FFTMacProjector.H
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
#ifndef AMREX_FFT_MAC_PROJECTOR_H_
#define AMREX_FFT_MAC_PROJECTOR_H_

#include <AMReX_FFT_Poisson.H>
#include <AMReX_LO_BCTYPES.H>

namespace Hydro {

class FFTMacProjector
{
public:

FFTMacProjector (amrex::Geometry const& geom,
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> const& lobc,
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> const& hibc);

void setUMAC (amrex::Array<amrex::MultiFab*,AMREX_SPACEDIM> const& umac);

void project ();

private:
amrex::Geometry m_geom;
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> m_lobc;
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> m_hibc;
amrex::FFT::Poisson<amrex::MultiFab> m_fft;
amrex::Array<amrex::MultiFab*,AMREX_SPACEDIM> m_umac;
};

}
#endif
114 changes: 114 additions & 0 deletions Projections/hydro_FFTMacProjector.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
#include <hydro_FFTMacProjector.H>
#include <AMReX_MultiFabUtil.H>

using namespace amrex;

namespace Hydro {

namespace {
Array<std::pair<FFT::Boundary,FFT::Boundary>,AMREX_SPACEDIM>
make_fft_bc (amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> const& lobc,
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> const& hibc)
{
Array<std::pair<FFT::Boundary,FFT::Boundary>,AMREX_SPACEDIM> fft_bc;
for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
if ((lobc[idim] == LinOpBCType::Periodic &&
hibc[idim] != LinOpBCType::Periodic) ||
(lobc[idim] != LinOpBCType::Periodic &&
hibc[idim] == LinOpBCType::Periodic))
{
amrex::Abort("FFTMacProjector: wrong periodic BC type");
}
for (int face = 0; face < 2; ++face) {
auto& lbc = (face == 0) ? lobc[idim] : hibc[idim];
auto& fbc = (face == 0) ? fft_bc[idim].first : fft_bc[idim].second;
switch (lbc)
{
case LinOpBCType::Periodic:
fbc = FFT::Boundary::periodic;
break;
case LinOpBCType::Dirichlet:
fbc = FFT::Boundary::odd;
break;
case LinOpBCType::Neumann:
fbc = FFT::Boundary::even;
break;
default:
amrex::Abort("FFTMacProjector: only Periodic, Dirichlet & Neumann are supported");
}
}
}
return fft_bc;
}
}

FFTMacProjector::FFTMacProjector (
amrex::Geometry const& geom,
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> const& lobc,
amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> const& hibc)
: m_geom(geom),
m_lobc(lobc),
m_hibc(hibc),
m_fft(geom, make_fft_bc(lobc,hibc))
{}

void FFTMacProjector::setUMAC (Array<MultiFab*,AMREX_SPACEDIM> const& umac)
{
m_umac = umac;
}

void FFTMacProjector::project ()
{
MultiFab rhs(amrex::convert(m_umac[0]->boxArray(), IntVect(0)),
m_umac[0]->DistributionMap(), 1, 0);
amrex::computeDivergence(rhs, GetArrOfConstPtrs(m_umac), m_geom);

AMREX_ALWAYS_ASSERT(m_geom.Domain().numPts() == rhs.boxArray().numPts());

bool has_dirichlet = false;
for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
has_dirichlet = has_dirichlet ||
m_lobc[idim] == LinOpBCType::Dirichlet ||
m_hibc[idim] == LinOpBCType::Dirichlet;
}
if (! has_dirichlet) {
auto rhosum = rhs.sum(0);
rhs.plus(-rhosum/m_geom.Domain().d_numPts(), 0, 1);
}

MultiFab phi(rhs.boxArray(), rhs.DistributionMap(), 1, 1);
m_fft.solve(phi, rhs);
phi.FillBoundary(m_geom.periodicity());

auto const& phima = phi.const_arrays();
Box const& domain = m_geom.growPeriodicDomain(1);
for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
auto const& uma = m_umac[idim]->arrays();
int domlo = domain.smallEnd(idim);
int domhi = domain.bigEnd(idim)+1;
auto const bclo = m_lobc[idim];
auto const bchi = m_hibc[idim];
Real const dxinv = m_geom.InvCellSize(idim);
ParallelFor(*m_umac[idim], [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
IntVect iv(AMREX_D_DECL(i,j,k));
IntVect miv = iv - IntVect::TheDimensionVector(idim);
auto const& u = uma[b];
auto const& p = phima[b];
if (iv[idim] == domlo) {
if (bclo == LinOpBCType::Dirichlet) {
u(i,j,k) -= Real(2.0)*dxinv*p(iv);
}
} else if (iv[idim] == domhi) {
if (bchi == LinOpBCType::Dirichlet) {
u(i,j,k) += Real(2.0)*dxinv*p(miv);
}
} else {
u(i,j,k) -= dxinv*(p(iv)-p(miv));
}
});
}
Gpu::streamSynchronize();
}

}
14 changes: 4 additions & 10 deletions Projections/hydro_MacProjector.H
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,14 @@
#include <AMReX_MLEBABecLap.H>
#endif

#ifdef AMREX_USE_FFT
#include <AMReX_FFT_Poisson.H>
#endif

#ifdef AMREX_USE_HYPRE
#include <AMReX_HypreMLABecLap.H>
#endif

#ifdef AMREX_USE_FFT
#include <hydro_FFTMacProjector.H>
#endif

namespace Hydro {

class MacProjector
Expand Down Expand Up @@ -216,12 +216,6 @@ private:
#ifdef AMREX_USE_HYPRE
std::unique_ptr<amrex::HypreMLABecLap> m_hypremlabeclap;
#endif

bool m_use_fft = false;
#ifdef AMREX_USE_FFT
std::unique_ptr<amrex::FFT::Poisson<amrex::MultiFab>> m_fft_poisson;
std::unique_ptr<amrex::FFT::PoissonHybrid<amrex::MultiFab>> m_fft_poisson_hybrid;
#endif
};

}
Expand Down
54 changes: 2 additions & 52 deletions Projections/hydro_MacProjector.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ MacProjector::MacProjector(
m_phi_loc(a_phi_loc),
m_divu_loc(a_divu_loc)
{
amrex::ignore_unused(m_divu_loc, m_beta_loc, m_phi_loc, m_umac_loc, m_use_fft);
amrex::ignore_unused(m_divu_loc, m_beta_loc, m_phi_loc, m_umac_loc);
}

MacProjector::MacProjector (const Vector<Array<MultiFab*,AMREX_SPACEDIM> >& a_umac,
Expand Down Expand Up @@ -259,13 +259,6 @@ MacProjector::setDomainBC (const Array<LinOpBCType,AMREX_SPACEDIM>& lobc,
#ifdef AMREX_USE_HYPRE
m_hypremlabeclap.reset();
#endif
#ifdef AMREX_USE_FFT
if (m_fft_poisson_hybrid) {
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(lobc.back() == LinOpBCType::Neumann &&
hibc.back() == LinOpBCType::Neumann,
"FFT::PoissonHybrid supports Neumann BC in z-direction only");
}
#endif
}


Expand Down Expand Up @@ -394,18 +387,6 @@ MacProjector::project (Real reltol, Real atol)
m_mlmg->prepareForFluxes(amrex::GetVecOfConstPtrs(m_phi));
}
} else
#endif
#ifdef AMREX_USE_FFT
if (m_fft_poisson || m_fft_poisson_hybrid) {
m_mlmg->makeSolvable(0,0,m_rhs[0]);
if (m_fft_poisson) {
m_fft_poisson->solve(m_phi[0], m_rhs[0]);
} else {
m_fft_poisson_hybrid->solve(m_phi[0], m_rhs[0]);
}
// The solve below should not do any iterations.
m_mlmg->solve(amrex::GetVecOfPtrs(m_phi), amrex::GetVecOfConstPtrs(m_rhs), reltol, atol);
} else
#endif
{
m_mlmg->solve(amrex::GetVecOfPtrs(m_phi), amrex::GetVecOfConstPtrs(m_rhs), reltol, atol);
Expand Down Expand Up @@ -584,27 +565,6 @@ void MacProjector::initProjector (Vector<BoxArray> const& a_grids,
}
auto const& dm = a_dmap;

#ifdef AMREX_USE_FFT
{
ParmParse pp("mac_proj");
pp.query("use_fft", m_use_fft);
}
if (m_use_fft) {
bool fft_supported = (nlevs == 1) && (a_overset_mask.empty()) &&
m_geom[0].Domain().numPts() == ba[0].numPts();
#if (AMREX_SPACEDIM == 3)
fft_supported = fft_supported && (m_geom[0].isPeriodic(0) &&
m_geom[0].isPeriodic(1));
#else
fft_supported = fft_supported && m_geom[0].isAllPeriodic();
#endif
if (!fft_supported) {
m_use_fft = false;
amrex::Warning("MacProjector: FFT not support");
}
}
#endif

m_rhs.resize(nlevs);
m_phi.resize(nlevs);
m_fluxes.resize(nlevs);
Expand All @@ -622,7 +582,7 @@ void MacProjector::initProjector (Vector<BoxArray> const& a_grids,
}
}

if (m_use_mlhypre || m_use_fft) { a_lpinfo.setMaxCoarseningLevel(0); }
if (m_use_mlhypre) { a_lpinfo.setMaxCoarseningLevel(0); }

if (a_overset_mask.empty()) {
m_poisson = std::make_unique<MLPoisson>(m_geom, ba, dm, a_lpinfo);
Expand All @@ -648,16 +608,6 @@ void MacProjector::initProjector (Vector<BoxArray> const& a_grids,
}
#endif

#ifdef AMREX_USE_FFT
if (m_use_fft) {
if (m_geom[0].isAllPeriodic()) {
m_fft_poisson = std::make_unique<FFT::Poisson<MultiFab>>(m_geom[0]);
} else {
m_fft_poisson_hybrid = std::make_unique<FFT::PoissonHybrid<MultiFab>>(m_geom[0]);
}
}
#endif

setOptions();

m_needs_init = false;
Expand Down
41 changes: 41 additions & 0 deletions Tests/FFT_Mac_Projection/GNUmakefile
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
AMREX_HOME ?= ../../../amrex
AMREX_HYDRO_HOME = ../..

USE_FFT = TRUE
USE_MPI = TRUE
USE_OMP = FALSE

COMP = gnu

DIM = 3

DEBUG = FALSE

include $(AMREX_HOME)/Tools/GNUMake/Make.defs

include ./Make.package

Pdirs := AmrCore
Pdirs += Base
Pdirs += Boundary
Pdirs += FFT
Pdirs += LinearSolvers/MLMG

Ppack += $(foreach dir, $(Pdirs), $(AMREX_HOME)/Src/$(dir)/Make.package)
Ppack += $(AMREX_HYDRO_HOME)/Projections/Make.package

include $(Ppack)

Bdirs := AmrCore
Bdirs += Base
Bdirs += Boundary
Bdirs += FFT
Bdirs += LinearSolvers/MLMG

Blocs := $(foreach dir, $(Bdirs), $(AMREX_HOME)/Src/$(dir))
Blocs += $(AMREX_HYDRO_HOME)/Projections

INCLUDE_LOCATIONS += $(Blocs)
VPATH_LOCATIONS += $(Blocs)

include $(AMREX_HOME)/Tools/GNUMake/Make.rules
1 change: 1 addition & 0 deletions Tests/FFT_Mac_Projection/Make.package
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
CEXE_sources += main.cpp
Loading

0 comments on commit f89137f

Please sign in to comment.