Implementing sparse matrix representation of Jacobian and using PETSc preconditioners (BLAST-WarpX#6184)

Implemented a `MatrixPC` preconditioner that constructs the sparse
matrix representing the Jacobian for the implicit/semi-implicit time
integrators. Currently, this PC provides no native matrix solvers; it
relies on PETSc's preconditioners and matrix solvers. In the PETSc
interfaces for the JFNK algorithm, added the option of using
matrix-based preconditioners.

In this PR, the matrix comprises the Maxwell's equation terms
(discretized curl-curl) and the diagonal element of the mass matrix (for
the current term).

Other changes:
+ Implemented a wrapper for PETSc's SNES module to solve nonlinear
systems; it is an alternative to the native Newton solver.
+ Streamlined the PETSc interface by creating a `PETScSolver_impl` base
class that contains everything common to the linear and nonlinear
solvers. `KSP_impl` and `SNES_impl` now inherit from it.

Minor changes:
+ Moved PETSc initialization and finalizing calls to `main.cpp` to allow
for the use of command line inputs.
+ Simplified initializing the nonlinear solver based on user input by
removing `nlsolver_type_str` and directly using `AMREX_ENUM` types.

CI test will be added in a later PR.

Author: debog

Source/FieldSolver/ImplicitSolvers/ImplicitSolver.cpp

@@ -443,18 +443,29 @@ void ImplicitSolver::ComputeJfromMassMatrices (const bool  a_J_from_MM_only)
 void ImplicitSolver::parseNonlinearSolverParams ( const amrex::ParmParse&  pp )
 {
 
-    std::string nlsolver_type_str;
-    pp.get("nonlinear_solver", nlsolver_type_str);
+    pp.get("nonlinear_solver", m_nlsolver_type);
 
-    if (nlsolver_type_str=="picard") {
-        m_nlsolver_type = NonlinearSolverType::Picard;
+    if (m_nlsolver_type == NonlinearSolverType::picard) {
+
+        // Picard
         m_nlsolver = std::make_unique<PicardSolver<WarpXSolverVec,ImplicitSolver>>();
         m_max_particle_iterations = 1;
         m_particle_tolerance = 0.0;
+
     }
-    else if (nlsolver_type_str=="newton") {
-        m_nlsolver_type = NonlinearSolverType::Newton;
-        m_nlsolver = std::make_unique<NewtonSolver<WarpXSolverVec,ImplicitSolver>>();
+    else if (      (m_nlsolver_type == NonlinearSolverType::newton)
+                || (m_nlsolver_type == NonlinearSolverType::petsc_snes) ) {
+
+        // JFNK solvers
+        if (m_nlsolver_type == NonlinearSolverType::newton) {
+            m_nlsolver = std::make_unique<NewtonSolver<WarpXSolverVec,ImplicitSolver>>();
+        } else {
+#ifdef AMREX_USE_PETSC
+            m_nlsolver = std::make_unique<PETScSNES<WarpXSolverVec,ImplicitSolver>>();
+#else
+            WARPX_ABORT_WITH_MESSAGE("ImplicitSolver::parseNonlinearSolverParams(): must compile with PETSc to use petsc_snes (AMREX_USE_PETSC must be defined)");
+#endif
+        }
         pp.query("max_particle_iterations", m_max_particle_iterations);
         pp.query("particle_tolerance", m_particle_tolerance);
         pp.query("particle_suborbits", m_particle_suborbits);
@@ -846,7 +857,7 @@ void ImplicitSolver::PrintBaseImplicitSolverParameters () const
     amrex::Print() << "use particle suborbits:              " << (m_particle_suborbits ? "true":"false") << "\n";
     amrex::Print() << "print unconverged particle details:  " << (m_print_unconverged_particle_details ? "true":"false") << "\n";
     amrex::Print() << "Nonlinear solver type:               " << amrex::getEnumNameString(m_nlsolver_type) << "\n";
-    if (m_nlsolver_type==NonlinearSolverType::Newton) {
+    if (m_nlsolver_type==NonlinearSolverType::newton) {
         amrex::Print() << "use mass matrices:                   " << (m_use_mass_matrices ? "true":"false") << "\n";
         if (m_use_mass_matrices) {
             amrex::Print() << "    for jacobian calc:   " << (m_use_mass_matrices_jacobian ? "true":"false") << "\n";

Source/FieldSolver/ImplicitSolvers/WarpXSolverDOF.H

@@ -23,8 +23,10 @@ class WarpX;
  */
 struct WarpXSolverDOF
 {
-    ablastr::fields::MultiLevelVectorField m_array;
-    ablastr::fields::MultiLevelScalarField m_scalar;
+    amrex::Vector<std::array<std::unique_ptr<amrex::iMultiFab>,3>> m_array;
+    amrex::Vector<std::array<std::unique_ptr<amrex::iMultiFab>,3>> m_array_lhs;
+    amrex::Vector<std::unique_ptr<amrex::iMultiFab>> m_scalar;
+    amrex::Vector<std::unique_ptr<amrex::iMultiFab>> m_scalar_lhs;
 
     warpx::fields::FieldType m_array_type = warpx::fields::FieldType::None;
     warpx::fields::FieldType m_scalar_type = warpx::fields::FieldType::None;

Source/FieldSolver/ImplicitSolvers/WarpXSolverDOF.cpp

@@ -41,6 +41,8 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
 
     m_array.resize(a_num_amr_levels);
     m_scalar.resize(a_num_amr_levels);
+    m_array_lhs.resize(a_num_amr_levels);
+    m_scalar_lhs.resize(a_num_amr_levels);
 
     amrex::Long offset = 0;
     m_nDoFs_l = 0;
@@ -56,10 +58,10 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
             const ablastr::fields::VectorField this_array = a_WarpX->m_fields.get_alldirs(a_vector_type_name, lev);
             for (int n = 0; n < 3; n++) {
                 auto ncomp = this_array[n]->nComp();
-                m_array[lev][n] = new amrex::MultiFab( this_array[n]->boxArray(),
-                                                                this_array[n]->DistributionMap(),
-                                                                2*ncomp, // {local, global} for each comp
-                                                                amrex::IntVect::TheUnitVector() );
+                m_array[lev][n] = std::make_unique<amrex::iMultiFab>(this_array[n]->boxArray(),
+                                                                     this_array[n]->DistributionMap(),
+                                                                     2*ncomp, // {local, global} for each comp
+                                                                     this_array[n]->nGrowVect() );
                 m_nDoFs_g += this_array[n]->boxArray().numPts()*ncomp;
 
                 m_array[lev][n]->setVal(-1.0);
@@ -70,9 +72,10 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
                     ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                     {
                         for (int v = 0; v < ncomp; v++) {
-                            dof_arr(i,j,k,2*v) = (amrex::Real) bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))) * ncomp
-                                                 + (amrex::Real) offset_mf
-                                                 + (amrex::Real) offset;
+                            dof_arr(i,j,k,2*v) = bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))) * ncomp
+                                                 + v
+                                                 + offset_mf
+                                                 + offset;
                         }
                     });
                     offset_mf += bx.numPts()*ncomp;
@@ -94,10 +97,10 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
         for (int lev = 0; lev < a_num_amr_levels; ++lev) {
             const amrex::MultiFab* this_mf = a_WarpX->m_fields.get(a_scalar_type_name,lev);
             auto ncomp = this_mf->nComp();
-            m_scalar[lev] = new amrex::MultiFab( this_mf->boxArray(),
-                                                          this_mf->DistributionMap(),
-                                                          2*ncomp, // {local, global} for each comp
-                                                          amrex::IntVect::TheUnitVector() );
+            m_scalar[lev] = std::make_unique<amrex::iMultiFab>(this_mf->boxArray(),
+                                                               this_mf->DistributionMap(),
+                                                               2*ncomp, // {local, global} for each comp
+                                                               this_mf->nGrowVect() );
             m_nDoFs_g += this_mf->boxArray().numPts()*ncomp;
 
             m_scalar[lev]->setVal(-1.0);
@@ -108,9 +111,10 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
                 ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                 {
                     for (int v = 0; v < ncomp; v++) {
-                        dof_arr(i,j,k,2*v) = (amrex::Real) bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))) * ncomp
-                                             + (amrex::Real) offset_mf
-                                             + (amrex::Real) offset;
+                        dof_arr(i,j,k,2*v) = bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))) * ncomp
+                                             + v
+                                             + offset_mf
+                                             + offset;
                     }
                 });
                 offset_mf += bx.numPts()*ncomp;
@@ -147,7 +151,7 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
                     ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                     {
                         for (int v = 0; v < ncomp; v++) {
-                            dof_arr(i,j,k,2*v+1) = dof_arr(i,j,k,2*v) + (amrex::Real) offset_global;
+                            dof_arr(i,j,k,2*v+1) = dof_arr(i,j,k,2*v) + offset_global;
                         }
                     });
                 }
@@ -164,12 +168,39 @@ void WarpXSolverDOF::Define ( WarpX* const        a_WarpX,
                 ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                 {
                     for (int v = 0; v < ncomp; v++) {
-                        dof_arr(i,j,k,2*v+1) = dof_arr(i,j,k,2*v) + (amrex::Real) offset_global;
+                        dof_arr(i,j,k,2*v+1) = dof_arr(i,j,k,2*v) + offset_global;
                     }
                 });
             }
         }
     }
 
+    if (m_array_type != FieldType::None) {
+        for (int lev = 0; lev < a_num_amr_levels; ++lev) {
+            const auto& geom = a_WarpX->Geom(lev);
+            for (int n = 0; n < 3; n++) {
+                m_array_lhs[lev][n] = std::make_unique<amrex::iMultiFab>(m_array[lev][n]->boxArray(),
+                                                                         m_array[lev][n]->DistributionMap(),
+                                                                         m_array[lev][n]->nComp(),
+                                                                         0 );
+                amrex::iMultiFab::Copy(*m_array_lhs[lev][n], *m_array[lev][n], 0, 0, m_array[lev][n]->nComp(), 0);
+                m_array[lev][n]->FillBoundary(geom.periodicity());
+                // do NOT call FillBoundary() on m_array_lhs
+            }
+        }
+    }
+    if (m_scalar_type != FieldType::None) {
+        for (int lev = 0; lev < a_num_amr_levels; ++lev) {
+            m_scalar_lhs[lev] = std::make_unique<amrex::iMultiFab>(m_scalar[lev]->boxArray(),
+                                                                   m_scalar[lev]->DistributionMap(),
+                                                                   m_scalar[lev]->nComp(),
+                                                                   0 );
+            amrex::iMultiFab::Copy(*m_scalar_lhs[lev], *m_scalar[lev], 0, 0, m_scalar[lev]->nComp(), 0);
+            const auto& geom = a_WarpX->Geom(lev);
+            m_scalar[lev]->FillBoundary(geom.periodicity());
+            // do NOT call FillBoundary() on m_scalar_lhs
+        }
+    }
+
     amrex::Print() << "Defined DOF object for linear solves (total DOFs = " << m_nDoFs_g << ").\n";
 }

Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.H

@@ -310,10 +310,13 @@ public:
     void copyTo   (amrex::Real* const) const;
 
     // return WarpX pointer
-    inline auto getWarpX () const { return m_WarpX; }
+    [[nodiscard]] auto getWarpX () const { return m_WarpX; }
 
     // return the number of AMR levels
-    inline auto numAMRLevels () const { return m_num_amr_levels; }
+    [[nodiscard]] auto numAMRLevels () const { return m_num_amr_levels; }
+
+    // return DOFs object pointer
+    inline const auto& getDOFsObject () const { return m_dofs; }
 
 private:
 

Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.cpp

@@ -105,6 +105,7 @@ void WarpXSolverVec::Define ( WarpX*  a_WarpX,
     if (m_dofs == nullptr) {
         m_dofs = std::make_unique<WarpXSolverDOF>();
         m_dofs->Define(m_WarpX, m_num_amr_levels, m_vector_type_name, m_scalar_type_name);
+        amrex::ExecOnFinalize([p=&m_dofs] () { p->reset(); });
     }
 
     m_is_defined = true;
@@ -158,7 +159,7 @@ void WarpXSolverVec::copyFrom ( const amrex::Real* const a_arr)
                     ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                     {
                         for (int v = 0; v < ncomp; v++) {
-                            int dof = (int) dof_arr(i,j,k,2*v); // local
+                            int dof = dof_arr(i,j,k,2*v); // local
                             data_arr(i,j,k,v) = a_arr[dof];
                         }
                     });
@@ -174,7 +175,7 @@ void WarpXSolverVec::copyFrom ( const amrex::Real* const a_arr)
                 ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                 {
                     for (int v = 0; v < ncomp; v++) {
-                        int dof = (int) dof_arr(i,j,k,2*v); // local
+                        int dof = dof_arr(i,j,k,2*v); // local
                         data_arr(i,j,k,v) = a_arr[dof];
                     }
                 });
@@ -203,7 +204,7 @@ void WarpXSolverVec::copyTo ( amrex::Real* const a_arr) const
                     ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                     {
                         for (int v = 0; v < ncomp; v++) {
-                            int dof = (int) dof_arr(i,j,k,2*v); // local
+                            int dof = dof_arr(i,j,k,2*v); // local
                             a_arr[dof] = data_arr(i,j,k,v);
                         }
                     });
@@ -219,7 +220,7 @@ void WarpXSolverVec::copyTo ( amrex::Real* const a_arr) const
                 ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
                 {
                     for (int v = 0; v < ncomp; v++) {
-                        int dof = (int) dof_arr(i,j,k,2*v); // local
+                        int dof = dof_arr(i,j,k,2*v); // local
                         a_arr[dof] = data_arr(i,j,k,v);
                     }
                 });

Source/Initialization/WarpXInit.cpp

@@ -17,6 +17,11 @@
 #include <ablastr/parallelization/MPIInitHelpers.H>
 #include <ablastr/warn_manager/WarnManager.H>
 
+#ifdef AMREX_USE_PETSC
+#include <petscsys.h>
+#endif
+
+
 #include <optional>
 #include <string>
 
@@ -25,10 +30,19 @@ void warpx::initialization::initialize_external_libraries(int argc, char* argv[]
     ablastr::parallelization::mpi_init(argc, argv);
     warpx::initialization::amrex_init(argc, argv);
     ablastr::math::anyfft::setup();
+#ifdef AMREX_USE_PETSC
+    PETSC_COMM_WORLD = amrex::ParallelContext::CommunicatorSub();
+    PetscInitialize(&argc, &argv, nullptr, "WarpX with PETSc");
+    amrex::Print() << "Initialized PETSc.\n";
+#endif
 }
 
 void warpx::initialization::finalize_external_libraries ()
 {
+#ifdef AMREX_USE_PETSC
+    PetscFinalize();
+    amrex::Print() << "Finalized PETSc.\n";
+#endif
     ablastr::math::anyfft::cleanup();
     amrex::Finalize();
     ablastr::parallelization::mpi_finalize();

Source/NonlinearSolvers/JacobianFunctionMF.H

@@ -9,6 +9,7 @@
 
 #include "LinearFunction.H"
 #include "CurlCurlMLMGPC.H"
+#include "MatrixPC.H"
 #include "JacobiPC.H"
 #include "Utils/TextMsg.H"
 #include <AMReX_Config.H>
@@ -56,6 +57,18 @@ class JacobianFunctionMF : public LinearFunction<T,Ops>
         if (m_usePreCond) { m_preCond->Update(a_X); }
     }
 
+    inline
+    void getPCMatrix( amrex::Gpu::DeviceVector<int>& a_ridx_g,
+                      amrex::Gpu::DeviceVector<int>& a_nnz,
+                      amrex::Gpu::DeviceVector<int>& a_cidx_g,
+                      amrex::Gpu::DeviceVector<RT>&  a_aij,
+                      int& a_n, int& a_ncols_max ) override
+    {
+        AMREX_ALWAYS_ASSERT(m_usePreCond);
+        m_preCond->getPCMatrix(a_ridx_g, a_nnz, a_cidx_g, a_aij, a_n, a_ncols_max);
+    }
+
+
     T makeVecLHS () const override;
     T makeVecRHS () const override;
 
@@ -153,17 +166,9 @@ void JacobianFunctionMF<T,Ops>::define ( const T& a_U,
             m_preCond = std::make_unique<CurlCurlMLMGPC<T,Ops>>();
         } else if (m_pc_type == PreconditionerType::pc_jacobi) {
             m_preCond = std::make_unique<JacobiPC<T,Ops>>();
-        } else if (m_pc_type == PreconditionerType::pc_petsc) {
-#ifdef AMREX_USE_PETSC
-            WARPX_ABORT_WITH_MESSAGE("JacobianFunctionMF::Define(): pc_petsc not yet implemented");
-#else
-            WARPX_ABORT_WITH_MESSAGE("JacobianFunctionMF::Define(): must compile with PETSc to use pc_petsc (AMREX_USE_PETSC must be defined)");
-#endif
         } else {
-            std::stringstream convergenceMsg;
-            convergenceMsg << "JacobianFunctionMF::define(): " << amrex::getEnumNameString(m_pc_type)
-                            << " is not a valid preconditioner type.";
-            WARPX_ABORT_WITH_MESSAGE(convergenceMsg.str());
+            m_preCond = std::make_unique<MatrixPC<T,Ops>>();
+            m_preCond->setName(amrex::getEnumNameString(m_pc_type));
         }
         m_preCond->Define(a_U, a_ops);
     }

Source/NonlinearSolvers/LinearFunction.H

@@ -43,6 +43,13 @@ class LinearFunction
         //! update preconditioner
         virtual void updatePreCondMat ( const T&  a_X ) = 0;
 
+        //! get sparse matrix representation of preconditioner
+        virtual void getPCMatrix( amrex::Gpu::DeviceVector<int>&,
+                                  amrex::Gpu::DeviceVector<int>&,
+                                  amrex::Gpu::DeviceVector<int>&,
+                                  amrex::Gpu::DeviceVector<RT>&,
+                                  int&, int& ) = 0;
+
         //! create a new vector given a defined vector
         inline void create ( T& a_Z, const T& a_U )
         {