Merge pull request #143 from ajnonaka/GMRES

GMRES Tutorial

Docs/source/LinearSolvers_Tutorial.rst

@@ -26,6 +26,8 @@ Step-by-step instructions for configuring AMReX to use HYPRE for this example ar
 
 ``ABecLaplacian_F`` demonstrates how to solve with cell-centered data using the Fortran interfaces.
 
+``GMRES/Poisson`` demonstrates how to solve the Poisson equation using GMRES with Jacobi preconditioning.
+
 ``NodalPoisson`` demonstrates how to set up and solve a variable coefficient Poisson equation
 with the rhs and solution data on nodes.
 

ExampleCodes/LinearSolvers/GMRES/Poisson/AMReX_GMRES_Poisson.H

@@ -0,0 +1,215 @@
+#ifndef AMREX_GMRES_POISSON_H_
+#define AMREX_GMRES_POISSON_H_
+#include <AMReX_Config.H>
+
+#include <AMReX_GMRES.H>
+#include <utility>
+
+namespace amrex {
+
+/**
+ * \brief Solve Poisson's equation using amrex GMRES class
+ *
+ * Refer to comments in amrex/Src/LinearSolvers/AMReX_GMRES.H
+ * for details on function implementation requirements
+ */
+class GMRESPOISSON
+{
+public:
+    using RT = amrex::Real; // double or float
+    using GM = GMRES<MultiFab,GMRESPOISSON>;
+
+    explicit GMRESPOISSON (const BoxArray& ba, const DistributionMapping& dm, const Geometry& geom);
+
+    /**
+     * \brief Solve the linear system
+     *
+     * \param a_sol     unknowns, i.e., x in A x = b.
+     * \param a_rhs     RHS, i.e., b in A x = b.
+     * \param a_tol_rel relative tolerance.
+     * \param a_tol_abs absolute tolerance.
+     */
+    void solve (MultiFab& a_sol, MultiFab const& a_rhs, RT a_tol_rel, RT a_tol_abs);
+
+    //! Sets verbosity.
+    void setVerbose (int v) { m_gmres.setVerbose(v); }
+
+    //! Get the GMRES object.
+    GM& getGMRES () { return m_gmres; }
+
+    //! Make MultiFab without ghost cells
+    MultiFab makeVecRHS () const;
+
+    //! Make MultiFab with ghost cells and set ghost cells to zero
+    MultiFab makeVecLHS () const;
+
+    RT norm2 (MultiFab const& mf) const;
+
+    static void scale (MultiFab& mf, RT scale_factor);
+
+    RT dotProduct (MultiFab const& mf1, MultiFab const& mf2) const;
+
+    //! lhs = 0
+    static void setToZero (MultiFab& lhs);
+
+    //! lhs = rhs
+    static void assign (MultiFab& lhs, MultiFab const& rhs);
+
+    //! lhs += a*rhs
+    static void increment (MultiFab& lhs, MultiFab const& rhs, RT a);
+
+    //! lhs = a*rhs_a + b*rhs_b
+    static void linComb (MultiFab& lhs, RT a, MultiFab const& rhs_a, RT b, MultiFab const& rhs_b);
+
+    //! lhs = L(rhs)
+    void apply (MultiFab& lhs, MultiFab& rhs) const;
+
+    void precond (MultiFab& lhs, MultiFab const& rhs) const;
+
+    //! Control whether or not to use MLMG as preconditioner.
+    bool usePrecond (bool new_flag) { return std::exchange(m_use_precond, new_flag); }
+
+private:
+    GM m_gmres;
+    BoxArray m_ba;
+    DistributionMapping m_dm;
+    Geometry m_geom;
+    bool m_use_precond;
+};
+
+GMRESPOISSON::GMRESPOISSON (const BoxArray& ba, const DistributionMapping& dm, const Geometry& geom)
+    : m_ba(ba), m_dm(dm), m_geom(geom)
+{
+    m_gmres.define(*this);
+}
+
+auto GMRESPOISSON::makeVecRHS () const -> MultiFab
+{
+    return MultiFab(m_ba, m_dm, 1, 0);
+}
+
+auto GMRESPOISSON::makeVecLHS () const -> MultiFab
+{
+    return MultiFab(m_ba, m_dm, 1, 1);
+}
+
+auto GMRESPOISSON::norm2 (MultiFab const& mf) const -> RT
+{
+    return mf.norm2();
+}
+
+void GMRESPOISSON::scale (MultiFab& mf, RT scale_factor)
+{
+    mf.mult(scale_factor);
+}
+
+auto GMRESPOISSON::dotProduct (MultiFab const& mf1, MultiFab const& mf2) const -> RT
+{
+    return MultiFab::Dot(mf1,0,mf2,0,1,0);
+}
+
+void GMRESPOISSON::setToZero (MultiFab& lhs)
+{
+    lhs.setVal(0.);
+}
+
+void GMRESPOISSON::assign (MultiFab& lhs, MultiFab const& rhs)
+{
+    MultiFab::Copy(lhs,rhs,0,0,1,0);
+}
+
+void GMRESPOISSON::increment (MultiFab& lhs, MultiFab const& rhs, RT a)
+{
+    MultiFab::Saxpy(lhs,a,rhs,0,0,1,0);
+}
+
+void GMRESPOISSON::linComb (MultiFab& lhs, RT a, MultiFab const& rhs_a, RT b, MultiFab const& rhs_b)
+{
+    MultiFab::LinComb(lhs,a,rhs_a,0,b,rhs_b,0,0,1,0);
+}
+
+void GMRESPOISSON::apply (MultiFab& lhs, MultiFab& rhs) const
+{
+    // apply matrix to rhs for output lhs
+    rhs.FillBoundary(m_geom.periodicity());
+
+    const GpuArray<Real, AMREX_SPACEDIM> dx = m_geom.CellSizeArray();
+
+    for ( MFIter mfi(lhs,TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
+
+        const Box& bx = mfi.tilebox();
+
+        const Array4<const Real> & rhs_p = rhs.array(mfi);
+        const Array4<      Real> & lhs_p = lhs.array(mfi);
+
+        amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+        {
+            lhs_p(i,j,k) = ( rhs_p(i+1,j,k) - 2.*rhs_p(i,j,k) + rhs_p(i-1,j,k) ) / (dx[0]*dx[0])
+                         + ( rhs_p(i,j+1,k) - 2.*rhs_p(i,j,k) + rhs_p(i,j-1,k) ) / (dx[1]*dx[1])
+#if (AMREX_SPACEDIM == 3)
+                         + ( rhs_p(i,j,k+1) - 2.*rhs_p(i,j,k) + rhs_p(i,j,k-1) ) / (dx[2]*dx[2])
+#endif
+                ;
+        });
+
+    }
+
+
+}
+
+void GMRESPOISSON::precond (MultiFab& lhs, MultiFab const& rhs) const
+{
+    if (m_use_precond) {
+
+        // in the preconditioner we use right-preconditioning to solve
+        // lhs = P^inv(rhs), where P^inv approximates A^inv
+        // Here we use Jacobi iterations to represent P^inv with an initial guess of lhs=0
+
+        const GpuArray<Real, AMREX_SPACEDIM> dx = m_geom.CellSizeArray();
+
+        amrex::Real fac = 0.;
+        for (int d=0; d<AMREX_SPACEDIM; ++d) { fac -= 2./(dx[d]*dx[d]); }
+
+        MultiFab tmp(m_ba, m_dm, 1, 1);
+        auto const& tmp_ma = tmp.const_arrays();
+        auto const& rhs_ma = rhs.const_arrays();
+        auto const& lhs_ma = lhs.arrays();
+
+        lhs.setVal(0.);
+
+        const int niters = 8;
+        for (int iter = 0; iter < niters; ++iter) {
+
+            MultiFab::Copy(tmp, lhs, 0, 0, 1, 0);
+            tmp.FillBoundary(m_geom.periodicity());
+
+            ParallelFor(lhs, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
+            {
+                auto const& tmp_ptr = tmp_ma[b];
+                auto ax = ( tmp_ptr(i+1,j,k) - 2.*tmp_ptr(i,j,k) + tmp_ptr(i-1,j,k) ) / (dx[0]*dx[0])
+                        + ( tmp_ptr(i,j+1,k) - 2.*tmp_ptr(i,j,k) + tmp_ptr(i,j-1,k) ) / (dx[1]*dx[1])
+#if (AMREX_SPACEDIM == 3)
+                        + ( tmp_ptr(i,j,k+1) - 2.*tmp_ptr(i,j,k) + tmp_ptr(i,j,k-1) ) / (dx[2]*dx[2])
+#endif
+                    ;
+                auto res = rhs_ma[b](i,j,k) - ax;
+
+                lhs_ma[b](i,j,k) += res / fac * Real(2./3.); // 2/3: weighted jacobi
+
+            });
+            Gpu::streamSynchronize();
+
+        }
+    } else {
+        MultiFab::Copy(lhs,rhs,0,0,1,0);
+    }
+}
+
+void GMRESPOISSON::solve (MultiFab& a_sol, MultiFab const& a_rhs, RT a_tol_rel, RT a_tol_abs)
+{
+    m_gmres.solve(a_sol, a_rhs, a_tol_rel, a_tol_abs);
+}
+
+}
+
+#endif

ExampleCodes/LinearSolvers/GMRES/Poisson/GNUmakefile

@@ -0,0 +1,17 @@
+# AMREX_HOME defines the directory in which we will find all the AMReX code.
+AMREX_HOME ?= ../../../../../amrex
+
+DEBUG        = FALSE
+USE_MPI      = FALSE
+USE_OMP      = FALSE
+COMP         = gnu
+DIM          = 3
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+
+include ./Make.package
+
+include $(AMREX_HOME)/Src/Base/Make.package
+include $(AMREX_HOME)/Src/LinearSolvers/Make.package
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.rules

ExampleCodes/LinearSolvers/GMRES/Poisson/Make.package

@@ -0,0 +1,2 @@
+CEXE_sources += main.cpp
+CEXE_headers += AMReX_GMRES_Poisson.H

ExampleCodes/LinearSolvers/GMRES/Poisson/inputs

@@ -0,0 +1,4 @@
+n_cell = 32
+max_grid_size = 16
+
+use_precond = 1

ExampleCodes/LinearSolvers/GMRES/Poisson/main.cpp

@@ -0,0 +1,131 @@
+/*
+ * A simplified usage of the AMReX GMRES class
+ */
+
+#include <AMReX.H>
+#include <AMReX_PlotFileUtil.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_GMRES_Poisson.H>
+
+int main (int argc, char* argv[])
+{
+    amrex::Initialize(argc,argv);
+    {
+
+    // **********************************
+    // DECLARE SIMULATION PARAMETERS
+    // **********************************
+
+    // number of cells on each side of the domain
+    int n_cell;
+
+    // size of each box (or grid)
+    int max_grid_size;
+
+    // use preconditioner?
+    int use_precond;
+
+    // **********************************
+    // READ PARAMETER VALUES FROM INPUT DATA
+    // **********************************
+    // inputs parameters
+    {
+        // ParmParse is way of reading inputs from the inputs file
+        // pp.get means we require the inputs file to have it
+        // pp.query means we optionally need the inputs file to have it - but we must supply a default here
+        amrex::ParmParse pp;
+
+        // We need to get n_cell from the inputs file - this is the number of cells on each side of
+        //   a square (or cubic) domain.
+        pp.get("n_cell",n_cell);
+
+        // The domain is broken into boxes of size max_grid_size
+        pp.get("max_grid_size",max_grid_size);
+
+        use_precond = 0;
+        pp.query("use_precond",use_precond);
+    }
+
+    // **********************************
+    // DEFINE SIMULATION SETUP AND GEOMETRY
+    // **********************************
+
+    // make BoxArray and Geometry
+    // ba will contain a list of boxes that cover the domain
+    // geom contains information such as the physical domain size,
+    // number of points in the domain, and periodicity
+    amrex::BoxArray ba;
+    amrex::Geometry geom;
+
+    // define lower and upper indices
+    amrex::IntVect dom_lo(       0,        0,        0);
+    amrex::IntVect dom_hi(n_cell-1, n_cell-1, n_cell-1);
+
+    // Make a single box that is the entire domain
+    amrex::Box domain(dom_lo, dom_hi);
+
+    // Initialize the boxarray "ba" from the single box "domain"
+    ba.define(domain);
+
+    // Break up boxarray "ba" into chunks no larger than "max_grid_size" along a direction
+    ba.maxSize(max_grid_size);
+
+    // This defines the physical box, [0,1] in each direction.
+    amrex::RealBox real_box({ 0., 0., 0.},
+                            { 1., 1., 1.});
+
+    // periodic in all direction
+    amrex::Array<int,3> is_periodic{1,1,1};
+
+    // This defines a Geometry object
+    geom.define(domain, real_box, amrex::CoordSys::cartesian, is_periodic);
+
+    // How Boxes are distrubuted among MPI processes
+    amrex::DistributionMapping dm(ba);
+
+    // we allocate two phi multifabs; one will store the old state, the other the new.
+    amrex::MultiFab rhs(ba, dm, 1, 0);
+    amrex::MultiFab phi(ba, dm, 1, 1);
+
+    // **********************************
+    // INITIALIZE DATA LOOP
+    // **********************************
+
+    // loop over boxes
+    for (amrex::MFIter mfi(rhs); mfi.isValid(); ++mfi)
+    {
+        const amrex::Box& bx = mfi.validbox();
+
+        const amrex::Array4<amrex::Real>& rhs_p = rhs.array(mfi);
+
+        // fill rhs with random numbers
+        amrex::ParallelForRNG(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::RandomEngine const& engine) noexcept
+        {
+            rhs_p(i,j,k) = amrex::RandomNormal(0.,1.,engine);
+        });
+    }
+
+    // offset data so the rhs sums to zero
+    amrex::Real sum = rhs.sum();
+    amrex::Long npts = rhs.boxArray().numPts();
+    rhs.plus(-sum/npts,0,1);
+
+    WriteSingleLevelPlotfile("rhs", rhs, {"rhs"}, geom, 0., 0);
+
+    amrex::GMRESPOISSON gmres_poisson(ba,dm,geom);
+
+    // initial guess
+    phi.setVal(0.);
+
+    gmres_poisson.usePrecond(use_precond);
+    gmres_poisson.setVerbose(2);
+    gmres_poisson.solve(phi, rhs, 1.e-12, 0.);
+
+    WriteSingleLevelPlotfile("phi", phi, {"phi"}, geom, 0., 0);
+
+    }
+    amrex::Finalize();
+    return 0;
+}
+
+