Generalization of WarpXSolverVec class used by implicit solvers (ECP-…

…WarpX#5171)

* added warpx::fields::FieldType member to WarpXSolverVec.

* AMREX_ALWAYS ==> WARPX_ALWAYS

* Changed how WarpXSolverVec is defined. Explicit call to SetDotMask() is no longer required.

* refactoring.

* refactoring comments and how to Copy WarpXSolverVec.

* small comment fix.

* dotMask now owned by WarpX. Defined only when needed.

* clang tidy.

* adding Afield_dotMask.

* adding None as enum FieldType

* WarpXSolverVec can now be used with scalar field quantities.

* putting check for dotMask define inside SetDotMask().

* dont use namespace in header file.

* name change: field ==> array

* updating comments.

* updating comment.

* removed function.

* braces

* added isFieldArray() fun to warpx::fields::namespace.

* simplify logic for boolean return

* adding assert.

* adding assertSameType() function.

* adding assertIsDefined() function.

* both array_vec and scalar_vec types cant be FieldType::None

* restoring comment change in implicit solvers. moved to separate PR.

* additional revert to comments.

* one more revert to comments.

* fixing merge issue.

* reposition header file.

* adding ArrayFieldTypes[] to Fields.H

* using std::any_of()

* attempt to please clang tidy.

* Doxygen: Location of Fwd Declaration

Fix location of `WarpX`forward declaration.

---------

Co-authored-by: Axel Huebl <[email protected]>

Source/FieldSolver/Fields.H

@@ -7,10 +7,14 @@
 #ifndef WARPX_FIELDS_H_
 #define WARPX_FIELDS_H_
 
+#include <algorithm>
+#include <iterator>
+
 namespace warpx::fields
 {
     enum struct FieldType : int
     {
+        None,
         Efield_aux,
         Bfield_aux,
         Efield_fp,
@@ -37,6 +41,20 @@ namespace warpx::fields
         Efield_avg_cp,
         Bfield_avg_cp
     };
+
+    constexpr FieldType ArrayFieldTypes[] = {
+            FieldType::Efield_aux, FieldType::Bfield_aux, FieldType::Efield_fp, FieldType::Bfield_fp,
+            FieldType::current_fp, FieldType::current_fp_nodal, FieldType::vector_potential_fp,
+            FieldType::Efield_cp, FieldType::Bfield_cp, FieldType::current_cp,
+            FieldType::Efield_avg_fp, FieldType::Bfield_avg_fp, FieldType::Efield_avg_cp, FieldType::Bfield_avg_cp};
+
+    inline bool
+    isFieldArray (const FieldType field_type)
+    {
+        return std::any_of( std::begin(ArrayFieldTypes), std::end(ArrayFieldTypes),
+            [field_type](const FieldType& f) { return f == field_type; });
+    }
+
 }
 
 #endif //WARPX_FIELDS_H_

Source/FieldSolver/ImplicitSolvers/SemiImplicitEM.cpp

@@ -20,13 +20,8 @@ void SemiImplicitEM::Define ( WarpX*  a_WarpX )
     m_WarpX = a_WarpX;
 
     // Define E and Eold vectors
-    m_E.Define( m_WarpX->getMultiLevelField(FieldType::Efield_fp) );
-    m_Eold.Define( m_WarpX->getMultiLevelField(FieldType::Efield_fp) );
-
-    // Need to define the WarpXSolverVec owned dot_mask to do dot
-    // product correctly for linear and nonlinear solvers
-    const amrex::Vector<amrex::Geometry>& Geom = m_WarpX->Geom();
-    m_E.SetDotMask(Geom);
+    m_E.Define( m_WarpX, FieldType::Efield_fp );
+    m_Eold.Define( m_E );
 
     // Parse implicit solver parameters
     const amrex::ParmParse pp("implicit_evolve");
@@ -71,7 +66,7 @@ void SemiImplicitEM::OneStep ( amrex::Real  a_time,
     m_WarpX->SaveParticlesAtImplicitStepStart ( );
 
     // Save Eg at the start of the time step
-    m_Eold.Copy( m_WarpX->getMultiLevelField(FieldType::Efield_fp) );
+    m_Eold.Copy( FieldType::Efield_fp );
 
     // Advance WarpX owned Bfield_fp to t_{n+1/2}
     m_WarpX->EvolveB(a_dt, DtType::Full);

Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.cpp

@@ -21,13 +21,8 @@ void ThetaImplicitEM::Define ( WarpX* const  a_WarpX )
     m_WarpX = a_WarpX;
 
     // Define E and Eold vectors
-    m_E.Define( m_WarpX->getMultiLevelField(FieldType::Efield_fp) );
-    m_Eold.Define( m_WarpX->getMultiLevelField(FieldType::Efield_fp) );
-
-    // Need to define the WarpXSolverVec owned dot_mask to do dot
-    // product correctly for linear and nonlinear solvers
-    const amrex::Vector<amrex::Geometry>& Geom = m_WarpX->Geom();
-    m_E.SetDotMask(Geom);
+    m_E.Define( m_WarpX, FieldType::Efield_fp );
+    m_Eold.Define( m_E );
 
     // Define Bold MultiFab
     const int num_levels = 1;
@@ -92,7 +87,7 @@ void ThetaImplicitEM::OneStep ( const amrex::Real  a_time,
     m_WarpX->SaveParticlesAtImplicitStepStart ( );
 
     // Save Eg at the start of the time step
-    m_Eold.Copy( m_WarpX->getMultiLevelField(FieldType::Efield_fp) );
+    m_Eold.Copy( FieldType::Efield_fp );
 
     const int num_levels = static_cast<int>(m_Bold.size());
     for (int lev = 0; lev < num_levels; ++lev) {

Source/FieldSolver/ImplicitSolvers/WarpXImplicitOps.cpp

@@ -11,6 +11,7 @@
 #include "Diagnostics/ReducedDiags/MultiReducedDiags.H"
 #include "Evolve/WarpXDtType.H"
 #include "Evolve/WarpXPushType.H"
+#include "FieldSolver/Fields.H"
 #include "FieldSolver/FiniteDifferenceSolver/FiniteDifferenceSolver.H"
 #include "Parallelization/GuardCellManager.H"
 #include "Particles/MultiParticleContainer.H"
@@ -68,7 +69,11 @@ WarpX::ImplicitPreRHSOp ( amrex::Real  a_cur_time,
 void
 WarpX::SetElectricFieldAndApplyBCs ( const WarpXSolverVec&  a_E )
 {
-    const amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>, 3 > >& Evec = a_E.getVec();
+    WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
+        a_E.getArrayVecType()==warpx::fields::FieldType::Efield_fp,
+        "WarpX::SetElectricFieldAndApplyBCs() must be called with Efield_fp type");
+
+    const amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>, 3 > >& Evec = a_E.getArrayVec();
     amrex::MultiFab::Copy(*Efield_fp[0][0], *Evec[0][0], 0, 0, ncomps, Evec[0][0]->nGrowVect());
     amrex::MultiFab::Copy(*Efield_fp[0][1], *Evec[0][1], 0, 0, ncomps, Evec[0][1]->nGrowVect());
     amrex::MultiFab::Copy(*Efield_fp[0][2], *Evec[0][2], 0, 0, ncomps, Evec[0][2]->nGrowVect());
@@ -316,22 +321,26 @@ WarpX::ImplicitComputeRHSE (int lev, amrex::Real a_dt, WarpXSolverVec&  a_Erhs_v
 void
 WarpX::ImplicitComputeRHSE (int lev, PatchType patch_type, amrex::Real a_dt, WarpXSolverVec&  a_Erhs_vec)
 {
+    WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
+        a_Erhs_vec.getArrayVecType()==warpx::fields::FieldType::Efield_fp,
+        "WarpX::ImplicitComputeRHSE() must be called with Efield_fp type");
+
     // set RHS to zero value
-    a_Erhs_vec.getVec()[lev][0]->setVal(0.0);
-    a_Erhs_vec.getVec()[lev][1]->setVal(0.0);
-    a_Erhs_vec.getVec()[lev][2]->setVal(0.0);
+    a_Erhs_vec.getArrayVec()[lev][0]->setVal(0.0);
+    a_Erhs_vec.getArrayVec()[lev][1]->setVal(0.0);
+    a_Erhs_vec.getArrayVec()[lev][2]->setVal(0.0);
 
     // Compute Efield_rhs in regular cells by calling EvolveE. Because
     // a_Erhs_vec is set to zero above, calling EvolveE below results in
     // a_Erhs_vec storing only the RHS of the update equation. I.e.,
     // c^2*dt*(curl(B^{n+theta} - mu0*J^{n+1/2})
     if (patch_type == PatchType::fine) {
-        m_fdtd_solver_fp[lev]->EvolveE( a_Erhs_vec.getVec()[lev], Bfield_fp[lev],
+        m_fdtd_solver_fp[lev]->EvolveE( a_Erhs_vec.getArrayVec()[lev], Bfield_fp[lev],
                                         current_fp[lev], m_edge_lengths[lev],
                                         m_face_areas[lev], ECTRhofield[lev],
                                         F_fp[lev], lev, a_dt );
     } else {
-        m_fdtd_solver_cp[lev]->EvolveE( a_Erhs_vec.getVec()[lev], Bfield_cp[lev],
+        m_fdtd_solver_cp[lev]->EvolveE( a_Erhs_vec.getArrayVec()[lev], Bfield_cp[lev],
                                         current_cp[lev], m_edge_lengths[lev],
                                         m_face_areas[lev], ECTRhofield[lev],
                                         F_cp[lev], lev, a_dt );