Implicit Field Solve Preconditioner based on Curl-Curl Operator (#5286)

Implemented a preconditioner for the implicit E-field solve using the
AMReX curl-curl operator and the MLMG solver.
+ Introduced a `Preconditioner` base class that defines the action of a
preconditioner for the JFNK algorithm.
+ Implemented the `CurlCurlMLMGPC` that uses the multigrid solution for
the curl-curl operator (implemented in `AMReX`) to precondition the
E-field JFNK solve.

Other changes needed for this:
+ Partially implemented a mapping between WarpX field boundary types and
AMReX's linear operator boundary types.
+ Added some functionalities to `ImplicitSolver` class that allows
preconditioners to access `WarpX` info (like `Geometry`, boundaries,
etc).

Some premilinary wall times for:
```
Test: inputs_vandb_2d
  Grid: 160 X 160
  dt: 0.125/wpe = 2.22e-18 (dt_CFL = 7.84e-19 s, CFL = 2.83)
  Time iterations: 20

Solver parameters:
  newton.max_iterations = 10
  newton.relative_tolerance = 1.0e-12
  newton.absolute_tolerance = 0.0
  gmres.max_iterations = 1000
  gmres.relative_tolerance = 1.0e-8
  gmres.absolute_tolerance = 0.0

Avg GMRES iterations: ~3 (wPC), ~27 (noPC)
```

with `32^2` particles per cell:
```
Lassen (MPI + CUDA)
-------------------
  Box  GPU   Walltime (s)
             wPC       noPC
   1    1    2324.7    15004.1
   4    1    2306.8    14356.8
   4    4     758.9     3647.3

Dane (MPI + OMP)
----------------
  Box  CPU  Threads   Walltime (s)
                      wPC      noPC
   1    1      1      6709.3   43200.0*
   1    1      2      3279.1   22296.1
   1    1      4      1696.3   11613.2
   1    1      8      1085.0    6911.4
   1    1     16       724.3    4729.0
   4    1      1      5525.9   33288.8
  16    1      1      4419.4   28467.8
   4    4      1      1324.4    9121.1
  16   16      1       524.9    3658.8

* 43200.0 seconds is 12 hours (max job duration on Dane);
the simulation was almost done (started the 20th step).
```

with `10^2` particles per cell:
```
Lassen (MPI + CUDA)
-------------------
  Box  GPU   Walltime (s)
             wPC       noPC
   1    1    365.0     1443.5 
   4    1    254.1      927.8 
   4    4    133.1      301.5 

Dane (MPI + OMP)
----------------
  Box  CPU  Threads   Walltime (s)
                      wPC      noPC
   1    1      1      440.8    2360.5     
   1    1      2      241.7    1175.8 
   1    1      4      129.3     727.0 
   1    1      8       94.2     407.5 
   1    1     16       74.3     245.6 
   4    1      1      393.3    1932.5 
  16    1      1      337.6    1618.7 
   4    4      1       92.2     479.1 
  16   16      1       58.1     192.6 
```

---------

Co-authored-by: Edoardo Zoni <[email protected]>
Co-authored-by: Remi Lehe <[email protected]>
Co-authored-by: Justin Angus <[email protected]>
Co-authored-by: Weiqun Zhang <[email protected]>

Author: 5 people

Source/FieldSolver/ImplicitSolvers/CMakeLists.txt

@@ -2,6 +2,7 @@ foreach(D IN LISTS WarpX_DIMS)
     warpx_set_suffix_dims(SD ${D})
     target_sources(lib_${SD}
       PRIVATE
+        ImplicitSolver.cpp
         SemiImplicitEM.cpp
         ThetaImplicitEM.cpp
         WarpXImplicitOps.cpp

Source/FieldSolver/ImplicitSolvers/ImplicitSolver.H

@@ -1,4 +1,4 @@
-/* Copyright 2024 Justin Angus
+/* Copyright 2024 Justin Angus, Debojyoti Ghosh
  *
  * This file is part of WarpX.
  *
@@ -9,9 +9,11 @@
 
 #include "FieldSolver/ImplicitSolvers/WarpXSolverVec.H"
 #include "NonlinearSolvers/NonlinearSolverLibrary.H"
+#include "Utils/WarpXAlgorithmSelection.H"
 
 #include <AMReX_Array.H>
 #include <AMReX_REAL.H>
+#include <AMReX_LO_BCTYPES.H>
 
 /**
  * \brief Base class for implicit time solvers. The base functions are those
@@ -85,6 +87,16 @@ public:
                               int              a_nl_iter,
                               bool             a_from_jacobian ) = 0;
 
+    [[nodiscard]] virtual amrex::Real theta () const { return 1.0; }
+
+    [[nodiscard]] int numAMRLevels () const { return m_num_amr_levels; }
+
+    [[nodiscard]] const amrex::Geometry& GetGeometry (int) const;
+    [[nodiscard]] const amrex::Array<FieldBoundaryType,AMREX_SPACEDIM>& GetFieldBoundaryLo () const;
+    [[nodiscard]] const amrex::Array<FieldBoundaryType,AMREX_SPACEDIM>& GetFieldBoundaryHi () const;
+    [[nodiscard]] amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> GetLinOpBCLo () const;
+    [[nodiscard]] amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> GetLinOpBCHi () const;
+
 protected:
 
     /**
@@ -94,6 +106,11 @@ protected:
 
     bool m_is_defined = false;
 
+    /**
+     * \brief Number of AMR levels
+     */
+    int m_num_amr_levels = 1;
+
     /**
      * \brief Nonlinear solver type and object
      */
@@ -140,6 +157,11 @@ protected:
 
     }
 
+    /**
+     * \brief Convert from WarpX FieldBoundaryType to amrex::LinOpBCType
+     */
+    [[nodiscard]] amrex::Array<amrex::LinOpBCType,AMREX_SPACEDIM> convertFieldBCToLinOpBC ( const amrex::Array<FieldBoundaryType,AMREX_SPACEDIM>& ) const;
+
 };
 
 #endif

Source/FieldSolver/ImplicitSolvers/ImplicitSolver.cpp

@@ -0,0 +1,60 @@
+#include "ImplicitSolver.H"
+#include "WarpX.H"
+
+using namespace amrex;
+
+const Geometry& ImplicitSolver::GetGeometry (const int a_lvl) const
+{
+    AMREX_ASSERT((a_lvl >= 0) && (a_lvl < m_num_amr_levels));
+    return m_WarpX->Geom(a_lvl);
+}
+
+const Array<FieldBoundaryType,AMREX_SPACEDIM>& ImplicitSolver::GetFieldBoundaryLo () const
+{
+    return m_WarpX->GetFieldBoundaryLo();
+}
+
+const Array<FieldBoundaryType,AMREX_SPACEDIM>& ImplicitSolver::GetFieldBoundaryHi () const
+{
+    return m_WarpX->GetFieldBoundaryHi();
+}
+
+Array<LinOpBCType,AMREX_SPACEDIM> ImplicitSolver::GetLinOpBCLo () const
+{
+    return convertFieldBCToLinOpBC(m_WarpX->GetFieldBoundaryLo());
+}
+
+Array<LinOpBCType,AMREX_SPACEDIM> ImplicitSolver::GetLinOpBCHi () const
+{
+    return convertFieldBCToLinOpBC(m_WarpX->GetFieldBoundaryHi());
+}
+
+Array<LinOpBCType,AMREX_SPACEDIM> ImplicitSolver::convertFieldBCToLinOpBC (const Array<FieldBoundaryType,AMREX_SPACEDIM>& a_fbc) const
+{
+    Array<LinOpBCType, AMREX_SPACEDIM> lbc;
+    for (auto& bc : lbc) { bc = LinOpBCType::interior; }
+    for (int i = 0; i < AMREX_SPACEDIM; i++) {
+        if (a_fbc[i] == FieldBoundaryType::PML) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else if (a_fbc[i] == FieldBoundaryType::Periodic) {
+            lbc[i] = LinOpBCType::Periodic;
+        } else if (a_fbc[i] == FieldBoundaryType::PEC) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else if (a_fbc[i] == FieldBoundaryType::PMC) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else if (a_fbc[i] == FieldBoundaryType::Damped) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else if (a_fbc[i] == FieldBoundaryType::Absorbing_SilverMueller) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else if (a_fbc[i] == FieldBoundaryType::Neumann) {
+            lbc[i] = LinOpBCType::Neumann;
+        } else if (a_fbc[i] == FieldBoundaryType::None) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else if (a_fbc[i] == FieldBoundaryType::Open) {
+            WARPX_ABORT_WITH_MESSAGE("LinOpBCType not set for this FieldBoundaryType");
+        } else {
+            WARPX_ABORT_WITH_MESSAGE("Invalid value for FieldBoundaryType");
+        }
+    }
+    return lbc;
+}

Source/FieldSolver/ImplicitSolvers/Make.package

@@ -1,3 +1,4 @@
+CEXE_sources += ImplicitSolver.cpp
 CEXE_sources += SemiImplicitEM.cpp
 CEXE_sources += ThetaImplicitEM.cpp
 CEXE_sources += WarpXImplicitOps.cpp

Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.H

@@ -8,13 +8,12 @@
 #define THETA_IMPLICIT_EM_H_
 
 #include "FieldSolver/ImplicitSolvers/WarpXSolverVec.H"
+#include "ImplicitSolver.H"
 
 #include <AMReX_Array.H>
 #include <AMReX_MultiFab.H>
 #include <AMReX_REAL.H>
 
-#include "ImplicitSolver.H"
-
 /** @file
  *  Theta-implicit electromagnetic time solver class. This is a fully implicit
  *  algorithm where both the fields and particles are treated implicitly.
@@ -79,7 +78,7 @@ public:
                       int              a_nl_iter,
                       bool             a_from_jacobian ) override;
 
-    [[nodiscard]] amrex::Real theta () const { return m_theta; }
+    [[nodiscard]] amrex::Real theta () const override { return m_theta; }
 
 private:
 

Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.cpp

@@ -19,15 +19,15 @@ void ThetaImplicitEM::Define ( WarpX* const  a_WarpX )
 
     // Retain a pointer back to main WarpX class
     m_WarpX = a_WarpX;
+    m_num_amr_levels = 1;
 
     // Define E and Eold vectors
     m_E.Define( m_WarpX, "Efield_fp" );
     m_Eold.Define( m_E );
 
     // Define B_old MultiFabs
     using ablastr::fields::Direction;
-    const int num_levels = 1;
-    for (int lev = 0; lev < num_levels; ++lev) {
+    for (int lev = 0; lev < m_num_amr_levels; ++lev) {
         const auto& ba_Bx = m_WarpX->m_fields.get(FieldType::Bfield_fp, Direction{0}, lev)->boxArray();
         const auto& ba_By = m_WarpX->m_fields.get(FieldType::Bfield_fp, Direction{1}, lev)->boxArray();
         const auto& ba_Bz = m_WarpX->m_fields.get(FieldType::Bfield_fp, Direction{2}, lev)->boxArray();

Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.H

@@ -75,6 +75,7 @@ public:
     void Define ( const WarpXSolverVec&  a_solver_vec )
     {
         assertIsDefined( a_solver_vec );
+        m_num_amr_levels = a_solver_vec.m_num_amr_levels;
         Define( WarpXSolverVec::m_WarpX,
                 a_solver_vec.getVectorType(),
                 a_solver_vec.getScalarType() );
@@ -300,7 +301,7 @@ private:
     std::string m_scalar_type_name = "none";
 
     static constexpr int m_ncomp = 1;
-    static constexpr int m_num_amr_levels = 1;
+    int m_num_amr_levels = 1;
 
     inline static bool m_warpx_ptr_defined = false;
     inline static WarpX* m_WarpX = nullptr;

Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.cpp

@@ -34,6 +34,8 @@ void WarpXSolverVec::Define ( WarpX*  a_WarpX,
         m_warpx_ptr_defined = true;
     }
 
+    m_num_amr_levels = 1;
+
     m_vector_type_name = a_vector_type_name;
     m_scalar_type_name = a_scalar_type_name;