fix comments

randomize RHS Jacobi preconditioner (implemented but probably buggy as convergence slows down)
ajnonaka · Nov 27, 2024 · d5c0913 · d5c0913
1 parent 112d2f9
commit d5c0913
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 17 deletions.
diff --git a/ExampleCodes/LinearSolvers/GMRES/Poisson/AMReX_GMRES_Poisson.H b/ExampleCodes/LinearSolvers/GMRES/Poisson/AMReX_GMRES_Poisson.H
@@ -73,6 +73,9 @@ private:
     DistributionMapping m_dm;
     Geometry m_geom;
     bool m_use_precond = false;
+
+    // store the original RHS needed for Jacobi preconditioner
+    MultiFab orig_rhs;
 };
 
 template <typename MF>
@@ -144,7 +147,7 @@ void GMRESPOISSONT<MF>::apply (MF& lhs, MF& rhs) const
 
     const GpuArray<Real, AMREX_SPACEDIM> dx = m_geom.CellSizeArray();
 
-    for ( MFIter mfi(rhs,TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
+    for ( MFIter mfi(lhs,TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
 
         const Box& bx = mfi.tilebox();
 
@@ -166,11 +169,42 @@ void GMRESPOISSONT<MF>::apply (MF& lhs, MF& rhs) const
 
 }
 
+// applies preconditioner to rhs. If there is no preconditioner,
+// this function should do lhs = rhs.
 template <typename MF>
 void GMRESPOISSONT<MF>::precond (MF& lhs, MF const& rhs) const
 {
     if (m_use_precond) {
 
+        MultiFab rhs_tmp(m_ba,m_dm,1,1);
+        MultiFab::Copy(rhs_tmp,rhs,0,0,1,0);
+        rhs_tmp.FillBoundary(m_geom.periodicity());
+
+        const GpuArray<Real, AMREX_SPACEDIM> dx = m_geom.CellSizeArray();
+
+        amrex::Real fac = 1.;
+        for (int d=0; d<AMREX_SPACEDIM; ++d) fac += 2./(dx[d]*dx[d]);
+
+        for ( MFIter mfi(lhs,TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
+
+            const Box& bx = mfi.tilebox();
+
+            const Array4<const Real> & orig_rhs_p = orig_rhs.array(mfi);
+            const Array4<const Real> & rhs_p = rhs_tmp.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
+                    ;
+
+                lhs_p(i,j,k) = rhs_p(i,j,k) + (orig_rhs_p(i,j,k) - lhs_p(i,j,k) ) / fac;
+            });
+        }
     } else {
         MultiFab::Copy(lhs,rhs,0,0,1,0);
     }
@@ -179,6 +213,9 @@ void GMRESPOISSONT<MF>::precond (MF& lhs, MF const& rhs) const
 template <typename MF>
 void GMRESPOISSONT<MF>::solve (MF& a_sol, MF const& a_rhs, RT a_tol_rel, RT a_tol_abs)
 {
+    orig_rhs.define(m_ba,m_dm,1,0);
+    MultiFab::Copy(orig_rhs,a_rhs,0,0,1,0);
+
     m_gmres.solve(a_sol, a_rhs, a_tol_rel, a_tol_abs);
 }
 

diff --git a/ExampleCodes/LinearSolvers/GMRES/Poisson/main.cpp b/ExampleCodes/LinearSolvers/GMRES/Poisson/main.cpp
@@ -1,10 +1,7 @@
 /*
- * A simplified single file version of the HeatEquation_EX0_C exmaple.
- * This code is designed to be used with Demo_Tutorial.rst.
- *
+ * A simplified usage of the AMReX GMRES class
  */
 
-
 #include <AMReX.H>
 #include <AMReX_PlotFileUtil.H>
 #include <AMReX_ParmParse.H>
@@ -77,9 +74,6 @@ int main (int argc, char* argv[])
     // This defines a Geometry object
     geom.define(domain, real_box, amrex::CoordSys::cartesian, is_periodic);
 
-    // extract dx from the geometry object
-    amrex::GpuArray<amrex::Real,3> dx = geom.CellSizeArray();
-
     // How Boxes are distrubuted among MPI processes
     amrex::DistributionMapping dm(ba);
 
@@ -98,19 +92,18 @@ int main (int argc, char* argv[])
 
         const amrex::Array4<amrex::Real>& rhs_p = rhs.array(mfi);
 
-        // set rhs = 1 + e^(-(r-0.5)^2)
-        amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k)
+        // fill rhs with random numbers
+        amrex::ParallelForRNG(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::RandomEngine const& engine) noexcept
         {
-            // **********************************
-            // SET VALUES FOR EACH CELL
-            // **********************************
-            amrex::Real x = (i+0.5) * dx[0];
-            amrex::Real y = (j+0.5) * dx[1];
-            amrex::Real z = (k+0.5) * dx[2];
-            rhs_p(i,j,k) = sin(2.*M_PI*x) * sin(4.*M_PI*y) * sin(8.*M_PI*z);
+            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);