Skip to content

Commit 29f211e

Browse files
authored
Merge pull request #143 from ajnonaka/GMRES
GMRES Tutorial
2 parents 016d47c + cc09ed1 commit 29f211e

File tree

6 files changed

+371
-0
lines changed

6 files changed

+371
-0
lines changed

Docs/source/LinearSolvers_Tutorial.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ Step-by-step instructions for configuring AMReX to use HYPRE for this example ar
2626

2727
``ABecLaplacian_F`` demonstrates how to solve with cell-centered data using the Fortran interfaces.
2828

29+
``GMRES/Poisson`` demonstrates how to solve the Poisson equation using GMRES with Jacobi preconditioning.
30+
2931
``NodalPoisson`` demonstrates how to set up and solve a variable coefficient Poisson equation
3032
with the rhs and solution data on nodes.
3133

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
#ifndef AMREX_GMRES_POISSON_H_
2+
#define AMREX_GMRES_POISSON_H_
3+
#include <AMReX_Config.H>
4+
5+
#include <AMReX_GMRES.H>
6+
#include <utility>
7+
8+
namespace amrex {
9+
10+
/**
11+
* \brief Solve Poisson's equation using amrex GMRES class
12+
*
13+
* Refer to comments in amrex/Src/LinearSolvers/AMReX_GMRES.H
14+
* for details on function implementation requirements
15+
*/
16+
class GMRESPOISSON
17+
{
18+
public:
19+
using RT = amrex::Real; // double or float
20+
using GM = GMRES<MultiFab,GMRESPOISSON>;
21+
22+
explicit GMRESPOISSON (const BoxArray& ba, const DistributionMapping& dm, const Geometry& geom);
23+
24+
/**
25+
* \brief Solve the linear system
26+
*
27+
* \param a_sol unknowns, i.e., x in A x = b.
28+
* \param a_rhs RHS, i.e., b in A x = b.
29+
* \param a_tol_rel relative tolerance.
30+
* \param a_tol_abs absolute tolerance.
31+
*/
32+
void solve (MultiFab& a_sol, MultiFab const& a_rhs, RT a_tol_rel, RT a_tol_abs);
33+
34+
//! Sets verbosity.
35+
void setVerbose (int v) { m_gmres.setVerbose(v); }
36+
37+
//! Get the GMRES object.
38+
GM& getGMRES () { return m_gmres; }
39+
40+
//! Make MultiFab without ghost cells
41+
MultiFab makeVecRHS () const;
42+
43+
//! Make MultiFab with ghost cells and set ghost cells to zero
44+
MultiFab makeVecLHS () const;
45+
46+
RT norm2 (MultiFab const& mf) const;
47+
48+
static void scale (MultiFab& mf, RT scale_factor);
49+
50+
RT dotProduct (MultiFab const& mf1, MultiFab const& mf2) const;
51+
52+
//! lhs = 0
53+
static void setToZero (MultiFab& lhs);
54+
55+
//! lhs = rhs
56+
static void assign (MultiFab& lhs, MultiFab const& rhs);
57+
58+
//! lhs += a*rhs
59+
static void increment (MultiFab& lhs, MultiFab const& rhs, RT a);
60+
61+
//! lhs = a*rhs_a + b*rhs_b
62+
static void linComb (MultiFab& lhs, RT a, MultiFab const& rhs_a, RT b, MultiFab const& rhs_b);
63+
64+
//! lhs = L(rhs)
65+
void apply (MultiFab& lhs, MultiFab& rhs) const;
66+
67+
void precond (MultiFab& lhs, MultiFab const& rhs) const;
68+
69+
//! Control whether or not to use MLMG as preconditioner.
70+
bool usePrecond (bool new_flag) { return std::exchange(m_use_precond, new_flag); }
71+
72+
private:
73+
GM m_gmres;
74+
BoxArray m_ba;
75+
DistributionMapping m_dm;
76+
Geometry m_geom;
77+
bool m_use_precond;
78+
};
79+
80+
GMRESPOISSON::GMRESPOISSON (const BoxArray& ba, const DistributionMapping& dm, const Geometry& geom)
81+
: m_ba(ba), m_dm(dm), m_geom(geom)
82+
{
83+
m_gmres.define(*this);
84+
}
85+
86+
auto GMRESPOISSON::makeVecRHS () const -> MultiFab
87+
{
88+
return MultiFab(m_ba, m_dm, 1, 0);
89+
}
90+
91+
auto GMRESPOISSON::makeVecLHS () const -> MultiFab
92+
{
93+
return MultiFab(m_ba, m_dm, 1, 1);
94+
}
95+
96+
auto GMRESPOISSON::norm2 (MultiFab const& mf) const -> RT
97+
{
98+
return mf.norm2();
99+
}
100+
101+
void GMRESPOISSON::scale (MultiFab& mf, RT scale_factor)
102+
{
103+
mf.mult(scale_factor);
104+
}
105+
106+
auto GMRESPOISSON::dotProduct (MultiFab const& mf1, MultiFab const& mf2) const -> RT
107+
{
108+
return MultiFab::Dot(mf1,0,mf2,0,1,0);
109+
}
110+
111+
void GMRESPOISSON::setToZero (MultiFab& lhs)
112+
{
113+
lhs.setVal(0.);
114+
}
115+
116+
void GMRESPOISSON::assign (MultiFab& lhs, MultiFab const& rhs)
117+
{
118+
MultiFab::Copy(lhs,rhs,0,0,1,0);
119+
}
120+
121+
void GMRESPOISSON::increment (MultiFab& lhs, MultiFab const& rhs, RT a)
122+
{
123+
MultiFab::Saxpy(lhs,a,rhs,0,0,1,0);
124+
}
125+
126+
void GMRESPOISSON::linComb (MultiFab& lhs, RT a, MultiFab const& rhs_a, RT b, MultiFab const& rhs_b)
127+
{
128+
MultiFab::LinComb(lhs,a,rhs_a,0,b,rhs_b,0,0,1,0);
129+
}
130+
131+
void GMRESPOISSON::apply (MultiFab& lhs, MultiFab& rhs) const
132+
{
133+
// apply matrix to rhs for output lhs
134+
rhs.FillBoundary(m_geom.periodicity());
135+
136+
const GpuArray<Real, AMREX_SPACEDIM> dx = m_geom.CellSizeArray();
137+
138+
for ( MFIter mfi(lhs,TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
139+
140+
const Box& bx = mfi.tilebox();
141+
142+
const Array4<const Real> & rhs_p = rhs.array(mfi);
143+
const Array4< Real> & lhs_p = lhs.array(mfi);
144+
145+
amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
146+
{
147+
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])
148+
+ ( rhs_p(i,j+1,k) - 2.*rhs_p(i,j,k) + rhs_p(i,j-1,k) ) / (dx[1]*dx[1])
149+
#if (AMREX_SPACEDIM == 3)
150+
+ ( rhs_p(i,j,k+1) - 2.*rhs_p(i,j,k) + rhs_p(i,j,k-1) ) / (dx[2]*dx[2])
151+
#endif
152+
;
153+
});
154+
155+
}
156+
157+
158+
}
159+
160+
void GMRESPOISSON::precond (MultiFab& lhs, MultiFab const& rhs) const
161+
{
162+
if (m_use_precond) {
163+
164+
// in the preconditioner we use right-preconditioning to solve
165+
// lhs = P^inv(rhs), where P^inv approximates A^inv
166+
// Here we use Jacobi iterations to represent P^inv with an initial guess of lhs=0
167+
168+
const GpuArray<Real, AMREX_SPACEDIM> dx = m_geom.CellSizeArray();
169+
170+
amrex::Real fac = 0.;
171+
for (int d=0; d<AMREX_SPACEDIM; ++d) { fac -= 2./(dx[d]*dx[d]); }
172+
173+
MultiFab tmp(m_ba, m_dm, 1, 1);
174+
auto const& tmp_ma = tmp.const_arrays();
175+
auto const& rhs_ma = rhs.const_arrays();
176+
auto const& lhs_ma = lhs.arrays();
177+
178+
lhs.setVal(0.);
179+
180+
const int niters = 8;
181+
for (int iter = 0; iter < niters; ++iter) {
182+
183+
MultiFab::Copy(tmp, lhs, 0, 0, 1, 0);
184+
tmp.FillBoundary(m_geom.periodicity());
185+
186+
ParallelFor(lhs, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
187+
{
188+
auto const& tmp_ptr = tmp_ma[b];
189+
auto ax = ( tmp_ptr(i+1,j,k) - 2.*tmp_ptr(i,j,k) + tmp_ptr(i-1,j,k) ) / (dx[0]*dx[0])
190+
+ ( tmp_ptr(i,j+1,k) - 2.*tmp_ptr(i,j,k) + tmp_ptr(i,j-1,k) ) / (dx[1]*dx[1])
191+
#if (AMREX_SPACEDIM == 3)
192+
+ ( tmp_ptr(i,j,k+1) - 2.*tmp_ptr(i,j,k) + tmp_ptr(i,j,k-1) ) / (dx[2]*dx[2])
193+
#endif
194+
;
195+
auto res = rhs_ma[b](i,j,k) - ax;
196+
197+
lhs_ma[b](i,j,k) += res / fac * Real(2./3.); // 2/3: weighted jacobi
198+
199+
});
200+
Gpu::streamSynchronize();
201+
202+
}
203+
} else {
204+
MultiFab::Copy(lhs,rhs,0,0,1,0);
205+
}
206+
}
207+
208+
void GMRESPOISSON::solve (MultiFab& a_sol, MultiFab const& a_rhs, RT a_tol_rel, RT a_tol_abs)
209+
{
210+
m_gmres.solve(a_sol, a_rhs, a_tol_rel, a_tol_abs);
211+
}
212+
213+
}
214+
215+
#endif
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# AMREX_HOME defines the directory in which we will find all the AMReX code.
2+
AMREX_HOME ?= ../../../../../amrex
3+
4+
DEBUG = FALSE
5+
USE_MPI = FALSE
6+
USE_OMP = FALSE
7+
COMP = gnu
8+
DIM = 3
9+
10+
include $(AMREX_HOME)/Tools/GNUMake/Make.defs
11+
12+
include ./Make.package
13+
14+
include $(AMREX_HOME)/Src/Base/Make.package
15+
include $(AMREX_HOME)/Src/LinearSolvers/Make.package
16+
17+
include $(AMREX_HOME)/Tools/GNUMake/Make.rules
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
CEXE_sources += main.cpp
2+
CEXE_headers += AMReX_GMRES_Poisson.H
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
n_cell = 32
2+
max_grid_size = 16
3+
4+
use_precond = 1
Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
/*
2+
* A simplified usage of the AMReX GMRES class
3+
*/
4+
5+
#include <AMReX.H>
6+
#include <AMReX_PlotFileUtil.H>
7+
#include <AMReX_ParmParse.H>
8+
#include <AMReX_GMRES_Poisson.H>
9+
10+
int main (int argc, char* argv[])
11+
{
12+
amrex::Initialize(argc,argv);
13+
{
14+
15+
// **********************************
16+
// DECLARE SIMULATION PARAMETERS
17+
// **********************************
18+
19+
// number of cells on each side of the domain
20+
int n_cell;
21+
22+
// size of each box (or grid)
23+
int max_grid_size;
24+
25+
// use preconditioner?
26+
int use_precond;
27+
28+
// **********************************
29+
// READ PARAMETER VALUES FROM INPUT DATA
30+
// **********************************
31+
// inputs parameters
32+
{
33+
// ParmParse is way of reading inputs from the inputs file
34+
// pp.get means we require the inputs file to have it
35+
// pp.query means we optionally need the inputs file to have it - but we must supply a default here
36+
amrex::ParmParse pp;
37+
38+
// We need to get n_cell from the inputs file - this is the number of cells on each side of
39+
// a square (or cubic) domain.
40+
pp.get("n_cell",n_cell);
41+
42+
// The domain is broken into boxes of size max_grid_size
43+
pp.get("max_grid_size",max_grid_size);
44+
45+
use_precond = 0;
46+
pp.query("use_precond",use_precond);
47+
}
48+
49+
// **********************************
50+
// DEFINE SIMULATION SETUP AND GEOMETRY
51+
// **********************************
52+
53+
// make BoxArray and Geometry
54+
// ba will contain a list of boxes that cover the domain
55+
// geom contains information such as the physical domain size,
56+
// number of points in the domain, and periodicity
57+
amrex::BoxArray ba;
58+
amrex::Geometry geom;
59+
60+
// define lower and upper indices
61+
amrex::IntVect dom_lo( 0, 0, 0);
62+
amrex::IntVect dom_hi(n_cell-1, n_cell-1, n_cell-1);
63+
64+
// Make a single box that is the entire domain
65+
amrex::Box domain(dom_lo, dom_hi);
66+
67+
// Initialize the boxarray "ba" from the single box "domain"
68+
ba.define(domain);
69+
70+
// Break up boxarray "ba" into chunks no larger than "max_grid_size" along a direction
71+
ba.maxSize(max_grid_size);
72+
73+
// This defines the physical box, [0,1] in each direction.
74+
amrex::RealBox real_box({ 0., 0., 0.},
75+
{ 1., 1., 1.});
76+
77+
// periodic in all direction
78+
amrex::Array<int,3> is_periodic{1,1,1};
79+
80+
// This defines a Geometry object
81+
geom.define(domain, real_box, amrex::CoordSys::cartesian, is_periodic);
82+
83+
// How Boxes are distrubuted among MPI processes
84+
amrex::DistributionMapping dm(ba);
85+
86+
// we allocate two phi multifabs; one will store the old state, the other the new.
87+
amrex::MultiFab rhs(ba, dm, 1, 0);
88+
amrex::MultiFab phi(ba, dm, 1, 1);
89+
90+
// **********************************
91+
// INITIALIZE DATA LOOP
92+
// **********************************
93+
94+
// loop over boxes
95+
for (amrex::MFIter mfi(rhs); mfi.isValid(); ++mfi)
96+
{
97+
const amrex::Box& bx = mfi.validbox();
98+
99+
const amrex::Array4<amrex::Real>& rhs_p = rhs.array(mfi);
100+
101+
// fill rhs with random numbers
102+
amrex::ParallelForRNG(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::RandomEngine const& engine) noexcept
103+
{
104+
rhs_p(i,j,k) = amrex::RandomNormal(0.,1.,engine);
105+
});
106+
}
107+
108+
// offset data so the rhs sums to zero
109+
amrex::Real sum = rhs.sum();
110+
amrex::Long npts = rhs.boxArray().numPts();
111+
rhs.plus(-sum/npts,0,1);
112+
113+
WriteSingleLevelPlotfile("rhs", rhs, {"rhs"}, geom, 0., 0);
114+
115+
amrex::GMRESPOISSON gmres_poisson(ba,dm,geom);
116+
117+
// initial guess
118+
phi.setVal(0.);
119+
120+
gmres_poisson.usePrecond(use_precond);
121+
gmres_poisson.setVerbose(2);
122+
gmres_poisson.solve(phi, rhs, 1.e-12, 0.);
123+
124+
WriteSingleLevelPlotfile("phi", phi, {"phi"}, geom, 0., 0);
125+
126+
}
127+
amrex::Finalize();
128+
return 0;
129+
}
130+
131+

0 commit comments

Comments
 (0)