Extending preconditioners to handle a diagonal operator for E-field (BLAST-WarpX#5716)

- Extended the `CurlCurlMLMGPC` interface to work with non-scalar `beta`
coefficient. It can now be a `MultiFab` array (same size/layout as
E-Fields) with one component representing a diagonal matrix/operator on
the E-field.

- Implemented a `JacobiPC` preconditioner to solve the system `\sigma E
= b` using the Jacobi iterations. For now, `\sigma` is a diagonal matrix
(stored as a `MultiFab` array (same size/layout as E-Fields) with one
component. The solution is set as `E = b / \sigma`, so the iterations
are not yet implemented. The complete Jacobi method will be implemented
once non-diagonal elements of `\sigma` are available.

The preconditioners fetch a pointer to the matrix operator (`beta` or
`\sigma`) from the implicit time integrator object.

Author: debog

Source/FieldSolver/ImplicitSolvers/ImplicitSolver.H

@@ -66,6 +66,12 @@ public:
                            amrex::Real  a_dt,
                            int          a_step ) = 0;
 
+
+    /**
+     * \brief Return pointer to MultiFab array for mass matrix
+     */
+    [[nodiscard]] virtual const amrex::Vector<amrex::Array<amrex::MultiFab*,3>>* GetSigmaCoeff() const { return nullptr; }
+
     //
     // the following routines are called by the linear and nonlinear solvers
     //

Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.H

@@ -84,8 +84,21 @@ public:
                       int              a_nl_iter,
                       bool             a_from_jacobian ) override;
 
+    /**
+     * \brief Return pointer to MultiFab array for mass matrix
+     */
+    [[nodiscard]] const amrex::Vector<amrex::Array<amrex::MultiFab*,3>>* GetSigmaCoeff() const override
+    {
+        return &m_sigma_mfarrvec;
+    }
+
 private:
 
+    /**
+     * \brief Array of multifab pointers to mass matrix
+     */
+    amrex::Vector<amrex::Array<amrex::MultiFab*,3>> m_sigma_mfarrvec;
+
     /**
      * \brief Solver vectors to be used in the nonlinear solver to solve for the
      * electric field E. The main logic for determining which variables should be

Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.cpp

@@ -49,6 +49,26 @@ void ThetaImplicitEM::Define ( WarpX* const  a_WarpX )
     // Parse nonlinear solver parameters
     parseNonlinearSolverParams( pp );
 
+    // Define sigmaPC mutlifabs
+    using ablastr::fields::Direction;
+    for (int lev = 0; lev < m_num_amr_levels; ++lev) {
+        const auto& ba_Ex = m_WarpX->m_fields.get(FieldType::Efield_fp, Direction{0}, lev)->boxArray();
+        const auto& ba_Ey = m_WarpX->m_fields.get(FieldType::Efield_fp, Direction{1}, lev)->boxArray();
+        const auto& ba_Ez = m_WarpX->m_fields.get(FieldType::Efield_fp, Direction{2}, lev)->boxArray();
+        const auto& dm = m_WarpX->m_fields.get(FieldType::Efield_fp, Direction{0}, lev)->DistributionMap();
+        const amrex::IntVect ngb = m_WarpX->m_fields.get(FieldType::Efield_fp, Direction{0}, lev)->nGrowVect();
+        m_WarpX->m_fields.alloc_init(FieldType::sigmaPC, Direction{0}, lev, ba_Ex, dm, 1, ngb, 0.0_rt);
+        m_WarpX->m_fields.alloc_init(FieldType::sigmaPC, Direction{1}, lev, ba_Ey, dm, 1, ngb, 0.0_rt);
+        m_WarpX->m_fields.alloc_init(FieldType::sigmaPC, Direction{2}, lev, ba_Ez, dm, 1, ngb, 0.0_rt);
+    }
+
+    // Set the pointer to mass matrix MultiFab
+    for (int lev = 0; lev < m_num_amr_levels; ++lev) {
+        m_sigma_mfarrvec.push_back(m_WarpX->m_fields.get_alldirs(FieldType::sigmaPC, 0));
+        // setting m_sigma to 1.0 right now for testing
+        for (int dim = 0; dim < 3; dim++) { m_sigma_mfarrvec[lev][dim]->setVal(1.0); }
+    }
+
     // Define the nonlinear solver
     m_nlsolver->Define(m_E, this);
     m_is_defined = true;

Source/Fields.H

@@ -83,7 +83,8 @@ namespace warpx::fields
         Bfield_avg_cp,
         B_old, /**< Stores the value of B at the beginning of the timestep, for the implicit solver */
         ECTRhofield,
-        Venl
+        Venl,
+        sigmaPC
     );
 
     /** these are vector fields */
@@ -125,7 +126,8 @@ namespace warpx::fields
         FieldType::Bfield_avg_cp,
         FieldType::B_old,
         FieldType::ECTRhofield,
-        FieldType::Venl
+        FieldType::Venl,
+        FieldType::sigmaPC
     };
 
     /** Returns true if a FieldType represents a vector field */

Source/NonlinearSolvers/CurlCurlMLMGPC.H

@@ -75,12 +75,12 @@ class CurlCurlMLMGPC : public Preconditioner<T,Ops>
         /**
          * \brief Define the preconditioner
          */
-        void Define (const T&, Ops*) override;
+        void Define (const T&, Ops* const) override;
 
         /**
          * \brief Update the preconditioner
          */
-        void Update (const T&) override;
+        void Update (const T& a_U) override;
 
         /**
          * \brief Apply (solve) the preconditioner given a RHS
@@ -118,6 +118,8 @@ class CurlCurlMLMGPC : public Preconditioner<T,Ops>
         int m_max_iter = 10;
         int m_max_coarsening_level = 30;
 
+        bool m_beta_scalar = true;
+
         RT m_atol = 1.0e-16;
         RT m_rtol = 1.0e-4;
 
@@ -129,6 +131,8 @@ class CurlCurlMLMGPC : public Preconditioner<T,Ops>
         amrex::Vector<amrex::DistributionMapping> m_dmap;
         amrex::IntVect m_gv;
 
+        const amrex::Vector<amrex::Array<amrex::MultiFab*,3>>* m_bcoefs = nullptr;
+
         amrex::Array<amrex::LinOpBCType, AMREX_SPACEDIM> m_bc_lo;
         amrex::Array<amrex::LinOpBCType, AMREX_SPACEDIM> m_bc_hi;
 
@@ -153,6 +157,7 @@ void CurlCurlMLMGPC<T,Ops>::printParameters() const
     auto pc_name = getEnumNameString(PreconditionerType::pc_curl_curl_mlmg);
     Print() << pc_name << " verbose:              " << (m_verbose?"true":"false") << "\n";
     Print() << pc_name << " bottom verbose:       " << (m_bottom_verbose?"true":"false") << "\n";
+    Print() << pc_name << " beta coeff:           " << (m_beta_scalar?"scalar (1.0)":"MultiFab") << "\n";
     Print() << pc_name << " max iter:             " << m_max_iter << "\n";
     Print() << pc_name << " agglomeration:        " << m_agglomeration << "\n";
     Print() << pc_name << " consolidation:        " << m_consolidation << "\n";
@@ -203,18 +208,21 @@ void CurlCurlMLMGPC<T,Ops>::Define ( const T& a_U,
     // read preconditioner parameters
     readParameters();
 
+    // Get data vectors from a_U
+    auto& u_mfarrvec = a_U.getArrayVec();
+
     // create info object for curl-curl op
     m_info = std::make_unique<LPInfo>();
     m_info->setAgglomeration(m_agglomeration);
     m_info->setConsolidation(m_consolidation);
     m_info->setMaxCoarseningLevel(m_max_coarsening_level);
 
-    // Get data vectors from a_U
-    auto& u_mfarrvec = a_U.getArrayVec();
-
     // Set number of AMR levels and create geometry, grids, and
     // distribution mapping vectors.
     m_num_amr_levels = m_ops->numAMRLevels();
+    if (m_num_amr_levels > 1) {
+        WARPX_ABORT_WITH_MESSAGE("CurlCurlMLMGPC::Define(): m_num_amr_levels > 1");
+    }
     m_geom.resize(m_num_amr_levels);
     m_grids.resize(m_num_amr_levels);
     m_dmap.resize(m_num_amr_levels);
@@ -248,32 +256,54 @@ void CurlCurlMLMGPC<T,Ops>::Define ( const T& a_U,
         m_gmres_solver->setVerbose(static_cast<int>(m_verbose));
     }
 
+    m_bcoefs = m_ops->GetSigmaCoeff();
+    if (m_bcoefs != nullptr) { m_beta_scalar = false; }
+
     m_is_defined = true;
 }
 
 template <class T, class Ops>
 void CurlCurlMLMGPC<T,Ops>::Update (const T& a_U)
 {
+    using namespace amrex;
+
     WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
         IsDefined(),
         "CurlCurlMLMGPC::Update() called on undefined object" );
 
     // a_U is not needed for a linear operator
     amrex::ignore_unused(a_U);
 
-    // set the coefficients alpha and beta for curl-curl op
+    // set the alpha for curl-curl op
     const RT thetaDt = m_ops->GetTheta()*this->m_dt;
     const RT alpha = (thetaDt*PhysConst::c) * (thetaDt*PhysConst::c);
-    const RT beta = RT(1.0);
+    m_curl_curl->setScalars(alpha, RT(1.0));
 
-    m_curl_curl->setScalars(alpha, beta);
+    if (!m_beta_scalar) {
+        for (int n = 0; n < m_num_amr_levels; n++) {
+#if defined(WARPX_DIM_1D_Z)
+            // Missing dimensions is x,y in WarpX and y,z in AMReX
+            m_curl_curl->setBeta({Array<MultiFab const*,3>{ (*m_bcoefs)[n][2],
+                                                            (*m_bcoefs)[n][1],
+                                                            (*m_bcoefs)[n][0]}});
+#elif defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
+            // Missing dimension is y in WarpX and z in AMReX
+            m_curl_curl->setBeta({Array<MultiFab const*,3>{ (*m_bcoefs)[n][0],
+                                                            (*m_bcoefs)[n][2],
+                                                            (*m_bcoefs)[n][1]}});
+#elif defined(WARPX_DIM_3D)
+            m_curl_curl->setBeta({Array<MultiFab const*,3>{ (*m_bcoefs)[n][0],
+                                                            (*m_bcoefs)[n][1],
+                                                            (*m_bcoefs)[n][2]}});
+#endif
+        }
+    }
 
     if (m_verbose) {
         amrex::Print() << "Updating " << amrex::getEnumNameString(PreconditionerType::pc_curl_curl_mlmg)
                        << ": theta*dt = " << thetaDt << ", "
                        << " coefficients: "
-                       << "alpha = " << alpha << ", "
-                       << "beta = " << beta << "\n";
+                       << "alpha = " << alpha << "\n";
     }
 }