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 <59625522+EZoni@users.noreply.github.com> Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> Co-authored-by: Justin Angus <angus1@llnl.gov> Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>
ECP-WarpX · Nov 4, 2024 · c1cd7ab · c1cd7ab
1 parent 78cf034
commit c1cd7ab
Show file tree

Hide file tree

Showing 13 changed files with 615 additions and 24 deletions.
diff --git a/Source/FieldSolver/ImplicitSolvers/CMakeLists.txt b/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

diff --git a/Source/FieldSolver/ImplicitSolvers/ImplicitSolver.H b/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
diff --git a/Source/FieldSolver/ImplicitSolvers/ImplicitSolver.cpp b/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;
+}
diff --git a/Source/FieldSolver/ImplicitSolvers/Make.package b/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

diff --git a/Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.H b/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:
 

diff --git a/Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.cpp b/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();

diff --git a/Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.H b/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;

diff --git a/Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.cpp b/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;