Add FFT Based MAC Projector

Projections/CMakeLists.txt

@@ -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()

Projections/Make.package

@@ -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

Projections/hydro_FFTMacProjector.H

@@ -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

Projections/hydro_FFTMacProjector.cpp

@@ -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();
+}
+
+}

Projections/hydro_MacProjector.H

@@ -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
@@ -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
 };
 
 }

Projections/hydro_MacProjector.cpp

@@ -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,
@@ -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
 }
 
 
@@ -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);
@@ -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);
@@ -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);
@@ -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;

Tests/FFT_Mac_Projection/GNUmakefile

@@ -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

Tests/FFT_Mac_Projection/Make.package

@@ -0,0 +1 @@
+CEXE_sources += main.cpp