diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst index ced5eb4d374..3895cbf65a2 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -393,6 +393,30 @@ Overall simulation parameters - ``amrex_gmres.max_iterations`` (``int``, default: 1000) Maximum number of iterations. - ``amrex_gmres.relative_tolerance`` (``float``, default: 1.0e-4) Relative tolerance of the convergence. - ``amrex_gmres.absolute_tolerance`` (``float``, default: 0.0) Absolute tolerance of the convergence. + - ``amrex_gmres.pc_type`` (``string``, default: ``none``) Preconditioner applied inside the GMRES + iterations. The only supported options are ``none`` and ``pc_darwin_mlmg``, described below. + + - **Preconditioner options:** + Setting ``amrex_gmres.pc_type = pc_darwin_mlmg`` use the multi-grid algorithm + as a preconditioner within the GMRes iteration. Because the Darwin magnetoinductive equation + :math:`\nabla^4 Z + \nabla \times ( \chi(x) \nabla\times Z) = ...` is not well-adapted for multi-grid + (and because the preconditioner does not need to solve for the exact equation), here the multigrid + solver uses the approximate equation :math:`\nabla^2 ( \nabla^2 + \chi ) Z = ...`; this + is equivalent to the original magnetostatic equation if :math:`Z` is divergence-free and + `\chi` is a slowly varying function of space. In practice, two separate passes of multigrid are + used in the preconditioner, in order to invert the operators :math:`\nabla^2 + \chi` and `\nabla^2` + respectively. + + - ``pc_darwin_mlmg.verbose`` (``bool``, default: false) + - ``pc_darwin_mlmg.bottom_verbose`` (``bool``, default: false) + - ``pc_darwin_mlmg.agglomeration`` (``bool``, default: true) + - ``pc_darwin_mlmg.consolidation`` (``bool``, default: true) + - ``pc_darwin_mlmg.max_iter`` (``int``, default: 2) Fixed number of V-cycles per multigrid + solve. This is deliberately fixed, so that the preconditioner stays a fixed linear operator + over a GMRES solve (only true when solver tolerance is set to 0, as by default). + - ``pc_darwin_mlmg.max_coarsening_level`` (``int``, default: 30) + - ``pc_darwin_mlmg.relative_tolerance`` (``float``, default: 0) + - ``pc_darwin_mlmg.absolute_tolerance`` (``float``, default: 0) .. _param-electrostatic-pic: diff --git a/Examples/Tests/magnetized_plasma_modes/CMakeLists.txt b/Examples/Tests/magnetized_plasma_modes/CMakeLists.txt index bb90f0e5b27..7c6259a430c 100644 --- a/Examples/Tests/magnetized_plasma_modes/CMakeLists.txt +++ b/Examples/Tests/magnetized_plasma_modes/CMakeLists.txt @@ -35,7 +35,7 @@ add_warpx_test( test_2d_darwin_solver_em_modes_es_picmi # name 2 # dims 2 # nprocs - "inputs_test_em_modes_picmi.py --test --dim 2 --bdir z --darwin --include_es_solver" # inputs + "inputs_test_em_modes_picmi.py --test --dim 2 --bdir z --darwin --include_es_solver --use_preconditioner" # inputs "analysis.py --analyze_darwin_sim" # analysis "analysis_default_regression.py --path diags/field_diag000050" # checksum OFF # dependency diff --git a/Examples/Tests/magnetized_plasma_modes/inputs_test_em_modes_picmi.py b/Examples/Tests/magnetized_plasma_modes/inputs_test_em_modes_picmi.py index 851370009c6..e4540ba4b1b 100755 --- a/Examples/Tests/magnetized_plasma_modes/inputs_test_em_modes_picmi.py +++ b/Examples/Tests/magnetized_plasma_modes/inputs_test_em_modes_picmi.py @@ -95,6 +95,7 @@ def __init__( verbose, include_es_solver=False, use_rkf45=False, + use_preconditioner=False, ): """Get input parameters for the specific case desired.""" self.solver = solver @@ -104,6 +105,7 @@ def __init__( self.verbose = verbose or self.test self.include_es_solver = include_es_solver self.use_rkf45 = use_rkf45 + self.use_preconditioner = use_preconditioner # sanity check assert dim > 0 and dim < 4, f"{dim}-dimensions not a valid input" @@ -305,6 +307,11 @@ def setup_run(self): relative_tolerance=5e-5, max_iterations=2048, verbose_int=(2 if self.test else 0), + pc_type=( + picmi.DarwinMLMGPreconditioner() + if self.use_preconditioner + else None + ), ), ) if self.include_es_solver: @@ -541,6 +548,12 @@ def _record_average_fields(self): help="Ohm only: use adaptive RKF45 subcycling for the B-field update", action="store_true", ) +parser.add_argument( + "--use_preconditioner", + help="Darwin only: precondition the GMRES solve with the factored-Laplacian " + "multigrid preconditioner", + action="store_true", +) parser.add_argument( "-v", "--verbose", @@ -558,5 +571,6 @@ def _record_average_fields(self): verbose=args.verbose, include_es_solver=args.include_es_solver, use_rkf45=args.use_rkf45, + use_preconditioner=args.use_preconditioner, ) simulation.step() diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index b1a56b1b8e5..6f7afa32300 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -1654,6 +1654,13 @@ class GMRESLinearSolver(LinearSolverBase): absolute_tolerance: float, default=0. Absoluate tolerence of the convergence + + pc_type: preconditioner instance, optional + The preconditioner applied inside the GMRES iterations. This is only + used by solvers that drive GMRES directly rather than through a + nonlinear solver (currently the semi-implicit Darwin solver, which + supports an instance of DarwinMLMGPreconditioner); with a nonlinear + solver, pass the preconditioner to that solver instead. """ def __init__( @@ -1663,12 +1670,21 @@ def __init__( absolute_tolerance=None, relative_tolerance=None, max_iterations=None, + pc_type=None, ): self.verbose_int = verbose_int self.restart_length = restart_length self.absolute_tolerance = absolute_tolerance self.relative_tolerance = relative_tolerance self.max_iterations = max_iterations + self.pc_type = pc_type + + if pc_type is not None: + assert isinstance(pc_type, PreconditionerBase) + assert pc_type.supports_direct_gmres, ( + f"{type(pc_type).__name__} cannot be selected directly on the " + "GMRES solver; pass it to the nonlinear solver instead" + ) def linear_solver_initialize_inputs(self, nonlinear_solver=None): if nonlinear_solver is not None: @@ -1680,6 +1696,10 @@ def linear_solver_initialize_inputs(self, nonlinear_solver=None): amrex_gmres.relative_tolerance = self.relative_tolerance amrex_gmres.max_iterations = self.max_iterations + if self.pc_type is not None: + amrex_gmres.pc_type = self.pc_type.name + self.pc_type.preconditioner_type_initialize_inputs() + class PETScKSPLinearSolver(LinearSolverBase): """ @@ -1695,7 +1715,19 @@ def linear_solver_initialize_inputs(self, nonlinear_solver=None): class PreconditionerBase(picmistandard.base._ClassWithInit): - pass + # Name of the WarpX preconditioner type, set by subclasses. + name = None + + # Whether this preconditioner can be selected directly on a linear + # solver (rather than only via a nonlinear solver's Jacobian). + supports_direct_gmres = False + + def preconditioner_type_initialize_inputs(self, jacobian=None): + if jacobian is not None: + jacobian.pc_type = self.name + bucket = pywarpx.warpx.get_bucket(self.name) + for attr, value in vars(self).items(): + setattr(bucket, attr, value) class CurlCurlMLMGPreconditioner(PreconditionerBase): @@ -1727,6 +1759,8 @@ class CurlCurlMLMGPreconditioner(PreconditionerBase): Absoluate tolerence of the convergence """ + name = "pc_curl_curl_mlmg" + def __init__( self, verbose, @@ -1747,18 +1781,65 @@ def __init__( self.relative_tolerance = relative_tolerance self.absolute_tolerance = absolute_tolerance - def preconditioner_type_initialize_inputs(self, jacobian=None): - if jacobian is not None: - jacobian.pc_type = "pc_curl_curl_mlmg" - pc_curl_curl_mlmg = pywarpx.warpx.get_bucket("pc_curl_curl_mlmg") - pc_curl_curl_mlmg.verbose = self.verbose - pc_curl_curl_mlmg.bottom_verbose = self.bottom_verbose - pc_curl_curl_mlmg.agglomeration = self.agglomeration - pc_curl_curl_mlmg.consolidation = self.consolidation - pc_curl_curl_mlmg.max_iter = self.max_iter - pc_curl_curl_mlmg.max_coarsening_level = self.max_coarsening_level - pc_curl_curl_mlmg.relative_tolerance = self.relative_tolerance - pc_curl_curl_mlmg.absolute_tolerance = self.absolute_tolerance + +class DarwinMLMGPreconditioner(PreconditionerBase): + """ + Sets up the factored-Laplacian multigrid preconditioner for the + semi-implicit Darwin solver's GMRES iteration. Approximates the Darwin + field operator by its constant-susceptibility factorization + (-nabla^2)(-nabla^2 + chi) and applies it as two successive scalar + multigrid solves (Poisson then Helmholtz with the spatially varying + susceptibility) per vector component. + + Parameters + ---------- + verbose: bool, default=False + Whether there is verbose output from the solver + + bottom_verbose: bool, optional + Whether there is verbose output from the bottom solver + + agglomeration: bool, optional + + consolidation: bool, optional + + max_iter: int, default=2 + The fixed number of V-cycles used for each of the two multigrid + solves per component (fixed so the preconditioner is a fixed linear + operator across a GMRES solve) + + max_coarsening_level: int, optional + Maximum coarsening level + + relative_tolerance: float, optional + Relative tolerance of the convergence + + absolute_tolerance: float, optional + Absolute tolerance of the convergence + """ + + name = "pc_darwin_mlmg" + supports_direct_gmres = True + + def __init__( + self, + verbose=None, + bottom_verbose=None, + agglomeration=None, + consolidation=None, + max_iter=None, + max_coarsening_level=None, + relative_tolerance=None, + absolute_tolerance=None, + ): + self.verbose = verbose + self.bottom_verbose = bottom_verbose + self.agglomeration = agglomeration + self.consolidation = consolidation + self.max_iter = max_iter + self.max_coarsening_level = max_coarsening_level + self.relative_tolerance = relative_tolerance + self.absolute_tolerance = absolute_tolerance class JacobiPreconditioner(PreconditionerBase): @@ -1780,6 +1861,8 @@ class JacobiPreconditioner(PreconditionerBase): Absoluate tolerence of the convergence """ + name = "pc_jacobi" + def __init__( self, verbose, @@ -1792,15 +1875,6 @@ def __init__( self.relative_tolerance = relative_tolerance self.absolute_tolerance = absolute_tolerance - def preconditioner_type_initialize_inputs(self, jacobian=None): - if jacobian is not None: - jacobian.pc_type = "pc_jacobi" - pc_jacobi = pywarpx.warpx.get_bucket("pc_jacobi") - pc_jacobi.verbose = self.verbose - pc_jacobi.max_iter = self.max_iter - pc_jacobi.relative_tolerance = self.relative_tolerance - pc_jacobi.absolute_tolerance = self.absolute_tolerance - class PETScPreconditioner(PreconditionerBase): """ @@ -1827,6 +1901,8 @@ class PETScPreconditioner(PreconditionerBase): When type is "hypre" and hypre_type is "euclid" """ + name = "pc_petsc" + def __init__( self, type, @@ -1843,17 +1919,6 @@ def __init__( self.hypre_type = hypre_type self.euclid_factor_levels = euclid_factor_levels - def preconditioner_type_initialize_inputs(self, jacobian=None): - if jacobian is not None: - jacobian.pc_type = "pc_petsc" - pc_petsc = pywarpx.warpx.get_bucket("pc_petsc") - pc_petsc.type = self.type - pc_petsc.asm_overlap = self.asm_overlap - pc_petsc.sub_type = self.sub_type - pc_petsc.ilu_factor_levels = self.ilu_factor_levels - pc_petsc.hypre_type = self.hypre_type - pc_petsc.euclid_factor_levels = self.euclid_factor_levels - class NonlinearSolverBase(picmistandard.base._ClassWithInit): pass diff --git a/Regression/Checksum/benchmarks_json/test_2d_darwin_solver_em_modes_es_picmi.json b/Regression/Checksum/benchmarks_json/test_2d_darwin_solver_em_modes_es_picmi.json index fade661eedc..a3eb0c4ca69 100644 --- a/Regression/Checksum/benchmarks_json/test_2d_darwin_solver_em_modes_es_picmi.json +++ b/Regression/Checksum/benchmarks_json/test_2d_darwin_solver_em_modes_es_picmi.json @@ -1,26 +1,26 @@ { "electron": { - "particle_momentum_x": 6.34927151314996e-19, - "particle_momentum_y": 6.409422435076371e-19, - "particle_momentum_z": 6.66841263185746e-19, - "particle_position_x": 943.1663587753025, - "particle_position_y": 60355.016494036, + "particle_momentum_x": 6.348941729617574e-19, + "particle_momentum_y": 6.409610177870042e-19, + "particle_momentum_z": 6.668309310767983e-19, + "particle_position_x": 943.1657817267967, + "particle_position_y": 60355.2390525201, "particle_weight": 1.463940203059113e+17 }, "ions": { - "particle_momentum_x": 1.9313151265758452e-18, - "particle_momentum_y": 1.9266742336211786e-18, - "particle_momentum_z": 1.9233575044457035e-18, - "particle_position_x": 943.1875241816265, - "particle_position_y": 60358.90549056562, + "particle_momentum_x": 1.931296499833215e-18, + "particle_momentum_y": 1.926679490873102e-18, + "particle_momentum_z": 1.923355061015629e-18, + "particle_position_x": 943.1876258374203, + "particle_position_y": 60359.13572213739, "particle_weight": 1.463940203059113e+17 }, "lev=0": { - "Bx": 0.8518918672501272, - "By": 1.3910576389057014, + "Bx": 0.8516024083203164, + "By": 1.3912823741463232, "Bz": 307.2, - "Ex": 111758764.58429599, - "Ey": 33475017.00102582, - "Ez": 102229395.20342489 + "Ex": 111894288.60163084, + "Ey": 33775455.97767793, + "Ez": 102230555.84232913 } } diff --git a/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.H b/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.H index b2bc3a1a748..9f91d6d550d 100644 --- a/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.H +++ b/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.H @@ -15,6 +15,8 @@ #include +#include + class SemiImplicitDarwin; /** @@ -35,8 +37,8 @@ class SemiImplicitDarwin; * and it holds a pointer back to the solver for the fields, geometry and * mass matrices that the evaluation reads. * - * No preconditioner is currently implemented, so precond() applies the - * identity and define() rejects any preconditioner type other than `none`. + * The only preconditioner supported is `pc_darwin_mlmg` (see DarwinMLMGPC.H); + * with `none`, precond() applies the identity. */ class DarwinLinearFieldOperator final : public LinearFunction { @@ -60,16 +62,29 @@ public: */ void apply ( WarpXSolverVec& a_Ax, const WarpXSolverVec& a_x ) override; - /** \brief Apply the preconditioner. No preconditioner is implemented for - * the Darwin solver, so this applies the identity. */ + /** \brief Apply the preconditioner, i.e. approximately solve `P a_U = a_X`. + * With no preconditioner selected this applies the identity. */ inline void precond ( WarpXSolverVec& a_U, const WarpXSolverVec& a_X ) override { - a_U.Copy(a_X); + if (m_preCond) { m_preCond->Apply(a_U, a_X); } + else { a_U.Copy(a_X); } + } + + /** \brief Refresh the preconditioner from the current state of the solver + * (the freshly deposited mass matrices). A no-op without a preconditioner. */ + inline + void updatePreCondMat () override + { + if (m_preCond) { m_preCond->Update(); } } + /** \brief Print the preconditioner's parameters, if there is one. */ inline - void updatePreCondMat () override { } + void printParameters () const + { + if (m_preCond) { m_preCond->printParameters(); } + } inline void getPCMatrix ( amrex::Gpu::DeviceVector& a_ridx_g, @@ -93,14 +108,14 @@ public: * \param[in] a_U a defined solver vector with the layout of Z, used both * to make new vectors and to size the scratch space * \param[in] a_ops pointer back to the Darwin solver - * \param[in] a_pc_type preconditioner type; must be `none` + * \param[in] a_pc_type preconditioner type; `none` or `pc_darwin_mlmg` */ void define ( const WarpXSolverVec& a_U, SemiImplicitDarwin* a_ops, const PreconditionerType& a_pc_type ) override; [[nodiscard]] inline - PreconditionerType pcType () const override { return PreconditionerType::none; } + PreconditionerType pcType () const override { return m_pc_type; } private: @@ -112,6 +127,11 @@ private: /** \brief Pointer back to the Darwin solver */ SemiImplicitDarwin* m_ops = nullptr; + /** \brief Selected preconditioner type, and the preconditioner itself + * (null when no preconditioner is used). */ + PreconditionerType m_pc_type = PreconditionerType::none; + std::unique_ptr> m_preCond; + /** * \brief Scratch space used by apply(), allocated once in define() rather * than on every GMRES iteration since every iterate of Z shares the same diff --git a/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.cpp b/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.cpp index 020ea7b3f87..9d4f2f25f56 100644 --- a/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.cpp +++ b/Source/FieldSolver/ImplicitSolvers/DarwinLinearFieldOperator.cpp @@ -13,6 +13,8 @@ #include "Utils/TextMsg.H" #include "WarpX.H" +#include "NonlinearSolvers/DarwinMLMGPC.H" + #include using warpx::fields::FieldType; @@ -24,12 +26,20 @@ void DarwinLinearFieldOperator::define ( const WarpXSolverVec& a_U, BL_PROFILE("DarwinLinearFieldOperator::define()"); WARPX_ALWAYS_ASSERT_WITH_MESSAGE( - a_pc_type == PreconditionerType::none, - "DarwinLinearFieldOperator::define(): preconditioners are not supported"); + a_pc_type == PreconditionerType::none || + a_pc_type == PreconditionerType::pc_darwin_mlmg, + "DarwinLinearFieldOperator::define(): the only preconditioner supported " + "by the Darwin solver is pc_darwin_mlmg (or none for no preconditioning)"); m_R.Define(a_U); m_ops = a_ops; + m_pc_type = a_pc_type; + if (m_pc_type == PreconditionerType::pc_darwin_mlmg) { + m_preCond = std::make_unique>(); + m_preCond->Define(a_U, a_ops); + } + // Allocate the scratch space used by apply() once here (every iterate of // Z shares this same layout) rather than on every GMRES iteration. const auto& Zvec = a_U.getArrayVec(); diff --git a/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.H b/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.H index b7399586d99..89b4a9de1e4 100644 --- a/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.H +++ b/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.H @@ -79,6 +79,28 @@ public: void ApplyScaledMassMatrices ( ablastr::fields::MultiLevelVectorField& rhs, const ablastr::fields::MultiLevelVectorField& dA); + /** + * \brief Fill a cell-centered MultiFab with the scalar susceptibility + * chi(x): the (2 mu0/dt)-scaled row-sum of the diagonal mass-matrix + * bands, averaged over the three diagonal blocks and interpolated from + * their native (E-type) staggering to cell centers. This is the local + * chi scale the curl(chi curl Z) operator term applies to a uniform dA, + * used by the pc_darwin_mlmg preconditioner's Helmholtz factor. Must be + * called after the mass matrices have been deposited and synced for the + * current step. + * + * The full mass matrix is a 3x3 block tensor (one block per Jx/Jy/Jz + * row/column pair) at every point of each block's native E-type + * staggering, and each block is itself stored sparsely as several + * components representing the deposition-stencil bands rather than as + * a dense matrix. a_chi_cc collapses all of that to a single scalar per + * cell: it keeps only the 3 diagonal blocks (xx, yy, zz), sums each + * block's stencil-band components into a per-point row-sum, averages + * that row-sum over the 3 diagonal blocks, and interpolates the result + * from the staggered points onto the cell center. + */ + void ComputeScaledMassMatrixCC ( amrex::MultiFab& a_chi_cc ) const; + void PrepareVelocitiesForCurrentDeposition (); void AccumulateCurrentAndMassMatrices (); void CalculateSourceVector (); @@ -133,6 +155,13 @@ private: int m_linsol_restart_length = 30; amrex::Real m_linsol_atol = 0.; amrex::Real m_linsol_rtol = 1.0e-4; + + /** + * \brief Preconditioner for the GMRES solve (amrex_gmres.pc_type; off by + * default). Only PreconditionerType::pc_darwin_mlmg is supported: the + * factored-Laplacian MLMG preconditioner (see DarwinMLMGPC.H). + */ + PreconditionerType m_pc_type = PreconditionerType::none; }; #endif diff --git a/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.cpp b/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.cpp index 8bd9fbc5108..1d0e40ae4cf 100644 --- a/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.cpp +++ b/Source/FieldSolver/ImplicitSolvers/SemiImplicitDarwin.cpp @@ -70,11 +70,17 @@ void SemiImplicitDarwin::Define ( WarpX* a_WarpX, bool from_restart) pp_l.query("absolute_tolerance", m_linsol_atol); pp_l.query("relative_tolerance", m_linsol_rtol); pp_l.query("max_iterations", m_linsol_maxits); + pp_l.query("pc_type", m_pc_type); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + m_pc_type == PreconditionerType::none || + m_pc_type == PreconditionerType::pc_darwin_mlmg, + "The semi-implicit Darwin solver only supports pc_darwin_mlmg as " + "the GMRES preconditioner (amrex_gmres.pc_type)."); // Define the linear operator (this also allocates the scratch space it // uses to evaluate the operator on each GMRES iteration) m_linear_function = std::make_unique(); - m_linear_function->define(m_Z, this, PreconditionerType::none); + m_linear_function->define(m_Z, this, m_pc_type); // Define the linear solver if (m_linear_solver_type == LinearSolverType::amrex_gmres) { @@ -121,6 +127,8 @@ void SemiImplicitDarwin::PrintParameters () const amrex::Print() << "Linear solver (" << linsol_name << ") max iterations: " << m_linsol_maxits << "\n"; amrex::Print() << "Linear solver (" << linsol_name << ") relative tolerance: " << m_linsol_rtol << "\n"; amrex::Print() << "Linear solver (" << linsol_name << ") absolute tolerance: " << m_linsol_atol << "\n"; + amrex::Print() << "Linear solver (" << linsol_name << ") preconditioner: " << amrex::getEnumNameString(m_pc_type) << "\n"; + if (m_linear_function) { m_linear_function->printParameters(); } amrex::Print() << "-----------------------------------------------------------\n\n"; } @@ -183,6 +191,10 @@ int SemiImplicitDarwin::OneStep ( [[maybe_unused]] amrex::Real start_time, // i.e. fill m_source with `2 * laplacian(B) + 2 * mu_0 curl(J)` CalculateSourceVector(); + // Refresh the preconditioner from the freshly deposited mass matrices + // (no-op unless a preconditioner is enabled). + m_linear_function->updatePreCondMat(); + // Solve the magnetoinductive equation: // bilaplacian(Z) + curl(chi curl(Z)) = 2 * laplacian(B) + 2 * mu_0 curl(J) // where chi is the mass matrix scaled by 2 * mu_0 / dt (see @@ -579,3 +591,61 @@ void SemiImplicitDarwin::ApplyScaledMassMatrices ( rhs[lev][2]->FillBoundaryAndSync(m_WarpX->Geom(lev).periodicity()); } } + +void SemiImplicitDarwin::ComputeScaledMassMatrixCC ( amrex::MultiFab& a_chi_cc ) const +{ + BL_PROFILE("SemiImplicitDarwin::ComputeScaledMassMatrixCC()"); + + using ablastr::fields::Direction; + + const int lev = 0; + const amrex::MultiFab* Sdiag[3] = { + m_WarpX->m_fields.get(FieldType::MassMatrices_X, Direction{0}, lev), + m_WarpX->m_fields.get(FieldType::MassMatrices_Y, Direction{1}, lev), + m_WarpX->m_fields.get(FieldType::MassMatrices_Z, Direction{2}, lev)}; + + a_chi_cc.setVal(0.0); + + // Average over the three diagonal blocks and scale by the same 2 mu0/dt + // prefactor the operator applies to the mass-matrix product. + const amrex::Real fac = 2.0_rt * PhysConst::mu0 / (3.0_rt * m_dt); + + for (int d = 0; d < 3; ++d) { + const int nc = Sdiag[d]->nComp(); + const amrex::IntVect et = Sdiag[d]->ixType().toIntVect(); + const int e0 = et[0]; + const int e1 = (AMREX_SPACEDIM >= 2) ? et[1] : 0; + const int e2 = (AMREX_SPACEDIM >= 3) ? et[2] : 0; + // Each staggered point contributes with equal weight to the average + // onto the cell center (2 points per nodal dimension of the block). + const amrex::Real wt = + fac / static_cast((e0 + 1)*(e1 + 1)*(e2 + 1)); + +#ifdef AMREX_USE_OMP +#pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) +#endif + for (amrex::MFIter mfi(a_chi_cc, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi) + { + const amrex::Box& tbx = mfi.tilebox(); + amrex::Array4 const& chi = a_chi_cc.array(mfi); + amrex::Array4 const& S = Sdiag[d]->const_array(mfi); + amrex::ParallelFor(tbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + amrex::Real s = 0.0; + // Perform row-sum of mass matrix elements + for (int c = 0; c < nc; ++c) { + // Perform interpolation from `S`'s + // original staggering to the desired staggering + for (int kk = 0; kk <= e2; ++kk) { + for (int jj = 0; jj <= e1; ++jj) { + for (int ii = 0; ii <= e0; ++ii) { + s += S(i+ii,j+jj,k+kk,c); + } + } + } + } + chi(i,j,k) += wt*s; + }); + } + } +} diff --git a/Source/NonlinearSolvers/DarwinMLMGPC.H b/Source/NonlinearSolvers/DarwinMLMGPC.H new file mode 100644 index 00000000000..066852b2d3e --- /dev/null +++ b/Source/NonlinearSolvers/DarwinMLMGPC.H @@ -0,0 +1,359 @@ +/* Copyright 2026 The WarpX Community + * + * This file is part of WarpX. + * + * Authors: Roelof Groenewald (Realta Fusion) + * + * License: BSD-3-Clause-LBNL + */ +#ifndef DARWIN_MLMG_PC_H_ +#define DARWIN_MLMG_PC_H_ + +#include "Fields.H" +#include "Utils/TextMsg.H" +#include "Preconditioner.H" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +/** + * \brief Factored-Laplacian multigrid preconditioner for the semi-implicit + * Darwin solver's GMRES iteration. + * + * The Darwin field operator applied to the auxiliary variable Z is + * A(Z) = nabla^4(Z) + curl(chi curl(Z)), + * with chi = (2 mu0/dt) x the mass-matrix response. On the solenoidal + * subspace (div Z = 0) the plasma-response term is -chi nabla^2(Z), so in + * the constant-chi limit the operator factors exactly as + * A = (-nabla^2)(-nabla^2 + chi), + * and a spectrally equivalent preconditioner is two successive scalar + * elliptic solves: a Poisson solve followed by a Helmholtz solve with the + * spatially varying chi(x) as the MLABecLaplacian acoef. + * + * Z is B-staggered while MLABecLaplacian is cell-centered, so each vector + * component is collocated onto the cell-centered grid by index + * identification (node i -> cell i, a half-cell shift in the component's + * nodal dimension), solved, and mapped back the same way. Deliberately NOT + * averaged: pair-averaging multiplies the nodal-dimension Nyquist modes by + * cos(k dx/2) = 0, making the preconditioner exactly singular on those + * planes. + * + * The MLMG solves run a fixed number of V-cycles (max_iter) so the + * preconditioner is a fixed linear operator across a GMRES solve. + */ + +template +class DarwinMLMGPC : public Preconditioner +{ + public: + + using RT = typename T::value_type; + + /** + * \brief Default constructor + */ + DarwinMLMGPC () = default; + + /** + * \brief Default destructor + */ + ~DarwinMLMGPC () override = default; + + // Prohibit move and copy operations + DarwinMLMGPC (const DarwinMLMGPC&) = delete; + DarwinMLMGPC& operator= (const DarwinMLMGPC&) = delete; + DarwinMLMGPC (DarwinMLMGPC&&) noexcept = delete; + DarwinMLMGPC& operator= (DarwinMLMGPC&&) noexcept = delete; + + /** + * \brief Define the preconditioner + */ + void Define (const T&, Ops*) override; + + /** + * \brief Update the preconditioner: refresh chi(x) from the freshly + * deposited mass matrices (via Ops::ComputeScaledMassMatrixCC). + */ + void Update () override; + + /** + * \brief Apply (approximately solve) the preconditioner given a RHS: + * per Z component, collocate onto cell centers, solve + * (-nabla^2) y = b then (-nabla^2 + chi) x = y with fixed-cycle + * MLMG, and map back to the B-staggering. + */ + void Apply (T&, const T&) override; + + /** + * \brief Print parameters + */ + void printParameters () const override; + + /** + * \brief Check if this preconditioner has been defined. + */ + [[nodiscard]] inline bool IsDefined () const override { return m_is_defined; } + + protected: + + bool m_is_defined = false; + + bool m_verbose = false; + bool m_bottom_verbose = false; + bool m_agglomeration = true; + bool m_consolidation = true; + + int m_max_iter = 2; + int m_max_coarsening_level = 30; + + // Zero by default, so that the solve performs a fixed number of iterations, + // and the preconditioner is thus a fixed linear operator, as required by GMRES. + RT m_atol = 0.0; + RT m_rtol = 0.0; + + Ops* m_ops = nullptr; + + amrex::Geometry m_geom; + amrex::BoxArray m_grids; + amrex::DistributionMapping m_dmap; + + /** + * \brief Cell-centered susceptibility chi(x), refreshed in Update(). + */ + amrex::MultiFab m_chi_cc; + + /** + * \brief Cell-centered scratch shared by the per-component solves: + * RHS, intermediate (Poisson) solution, and final (Helmholtz) + * solution. The final solution carries one guard layer for the map + * back to the B-staggering (the top node plane in the nodal + * dimension reads the periodic image). + */ + amrex::MultiFab m_rhs_cc, m_mid_cc, m_sol_cc; + + std::unique_ptr m_info; + std::unique_ptr m_poisson; + std::unique_ptr m_helmholtz; + std::unique_ptr m_poisson_mg; + std::unique_ptr m_helmholtz_mg; + + /** + * \brief Read parameters + */ + void readParameters (); + + private: + +}; + +template +void DarwinMLMGPC::printParameters () const +{ + using namespace amrex; + auto pc_name = getEnumNameString(PreconditionerType::pc_darwin_mlmg); + Print() << pc_name << " verbose: " << (m_verbose?"true":"false") << "\n"; + Print() << pc_name << " bottom verbose: " << (m_bottom_verbose?"true":"false") << "\n"; + Print() << pc_name << " max iter (V-cycles): " << m_max_iter << "\n"; + Print() << pc_name << " agglomeration: " << m_agglomeration << "\n"; + Print() << pc_name << " consolidation: " << m_consolidation << "\n"; + Print() << pc_name << " max_coarsening_level: " << m_max_coarsening_level << "\n"; + Print() << pc_name << " absolute tolerance: " << m_atol << "\n"; + Print() << pc_name << " relative tolerance: " << m_rtol << "\n"; +} + +template +void DarwinMLMGPC::readParameters () +{ + const amrex::ParmParse pp(amrex::getEnumNameString(PreconditionerType::pc_darwin_mlmg)); + pp.query("verbose", m_verbose); + pp.query("bottom_verbose", m_bottom_verbose); + pp.query("max_iter", m_max_iter); + pp.query("agglomeration", m_agglomeration); + pp.query("consolidation", m_consolidation); + pp.query("max_coarsening_level", m_max_coarsening_level); + pp.query("absolute_tolerance", m_atol); + pp.query("relative_tolerance", m_rtol); +} + +template +void DarwinMLMGPC::Define ( const T& a_U, + Ops* const a_ops ) +{ + BL_PROFILE("DarwinMLMGPC::Define()"); + using namespace amrex; + + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + !IsDefined(), + "DarwinMLMGPC::Define() called on defined object" ); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + (a_ops != nullptr), + "DarwinMLMGPC::Define(): a_ops is nullptr" ); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + a_U.getArrayVecType()==warpx::fields::FieldType::Bfield_fp, + "DarwinMLMGPC::Define() must be called with a B-staggered solver vector"); + + m_ops = a_ops; + readParameters(); + + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + m_ops->numAMRLevels() == 1, + "DarwinMLMGPC::Define(): only a single AMR level is supported"); + + const auto& u_mfarrvec = a_U.getArrayVec(); + m_geom = m_ops->GetGeometry(0); + m_dmap = u_mfarrvec[0][0]->DistributionMap(); + m_grids = u_mfarrvec[0][0]->boxArray(); + m_grids.enclosedCells(); + + // The Darwin solver is restricted to fully periodic domains (see + // SemiImplicitDarwin::Define()), so the cell-centered operators are too. + Array bc_lo, bc_hi; + for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) { + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + m_geom.isPeriodic(idim), + "DarwinMLMGPC::Define(): only periodic boundaries are supported"); + bc_lo[idim] = LinOpBCType::Periodic; + bc_hi[idim] = LinOpBCType::Periodic; + } + + m_info = std::make_unique(); + m_info->setAgglomeration(m_agglomeration); + m_info->setConsolidation(m_consolidation); + m_info->setMaxCoarseningLevel(m_max_coarsening_level); + + // All three components share the same cell-centered grid, boundary + // conditions and coefficients, so one operator pair serves all of them. + m_poisson = std::make_unique( + Vector{m_geom}, Vector{m_grids}, + Vector{m_dmap}, *m_info); + m_poisson->setDomainBC(bc_lo, bc_hi); + m_poisson->setLevelBC(0, nullptr); + m_poisson->setScalars(RT(0.0), RT(1.0)); + m_poisson->setACoeffs(0, RT(0.0)); + m_poisson->setBCoeffs(0, RT(1.0)); + + m_helmholtz = std::make_unique( + Vector{m_geom}, Vector{m_grids}, + Vector{m_dmap}, *m_info); + m_helmholtz->setDomainBC(bc_lo, bc_hi); + m_helmholtz->setLevelBC(0, nullptr); + m_helmholtz->setScalars(RT(1.0), RT(1.0)); + m_helmholtz->setACoeffs(0, RT(0.0)); // chi(x) set in Update() + m_helmholtz->setBCoeffs(0, RT(1.0)); + + m_poisson_mg = std::make_unique(*m_poisson); + m_helmholtz_mg = std::make_unique(*m_helmholtz); + for (auto* mg : {m_poisson_mg.get(), m_helmholtz_mg.get()}) { + mg->setFixedIter(m_max_iter); + mg->setVerbose(static_cast(m_verbose)); + mg->setBottomVerbose(static_cast(m_bottom_verbose)); + } + + m_chi_cc.define(m_grids, m_dmap, 1, 0); + m_rhs_cc.define(m_grids, m_dmap, 1, 0); + m_mid_cc.define(m_grids, m_dmap, 1, 0); + m_sol_cc.define(m_grids, m_dmap, 1, 1); + + m_is_defined = true; +} + +template +void DarwinMLMGPC::Update () +{ + BL_PROFILE("DarwinMLMGPC::Update()"); + + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + IsDefined(), + "DarwinMLMGPC::Update() called on undefined object" ); + + m_ops->ComputeScaledMassMatrixCC(m_chi_cc); + m_helmholtz->setACoeffs(0, m_chi_cc); + + if (m_verbose) { + const amrex::Real chi_mean = + m_chi_cc.sum(0, false) / static_cast(m_chi_cc.boxArray().numPts()); + amrex::Print() << "Updating " + << amrex::getEnumNameString(PreconditionerType::pc_darwin_mlmg) + << ": mean chi = " << chi_mean << "\n"; + } +} + +template +void DarwinMLMGPC::Apply (T& a_x, const T& a_b) +{ + BL_PROFILE("DarwinMLMGPC::Apply()"); + using namespace amrex; + + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + IsDefined(), + "DarwinMLMGPC::Apply() called on undefined object" ); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + a_x.getArrayVecType()==warpx::fields::FieldType::Bfield_fp, + "DarwinMLMGPC::Apply() - a_x must be a B-staggered solver vector"); + WARPX_ALWAYS_ASSERT_WITH_MESSAGE( + a_b.getArrayVecType()==warpx::fields::FieldType::Bfield_fp, + "DarwinMLMGPC::Apply() - a_b must be a B-staggered solver vector"); + + const auto& b_mfarrvec = a_b.getArrayVec(); + const auto& x_mfarrvec = a_x.getArrayVec(); + const int lev = 0; + + for (int comp = 0; comp < 3; ++comp) { + + // Collocate the B-staggered RHS component onto the cell-centered + // grid by index identification (node i -> cell i; identity for a + // fully cell-centered component). The dropped top plane in the nodal + // dimension is the duplicate periodic image, so no information is + // lost. +#ifdef AMREX_USE_OMP +#pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) +#endif + for (MFIter mfi(m_rhs_cc, TilingIfNotGPU()); mfi.isValid(); ++mfi) + { + const Box& tbx = mfi.tilebox(); + Array4 const& rc = m_rhs_cc.array(mfi); + Array4 const& bs = b_mfarrvec[lev][comp]->const_array(mfi); + ParallelFor(tbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + rc(i,j,k) = bs(i,j,k); + }); + } + + // The two factored solves: Poisson then Helmholtz. Fixed V-cycle + // count (setFixedIter in Define) keeps this a fixed linear operator. + m_mid_cc.setVal(0.0); + m_poisson_mg->solve({&m_mid_cc}, {&m_rhs_cc}, m_rtol, m_atol); + m_sol_cc.setVal(0.0); + m_helmholtz_mg->solve({&m_sol_cc}, {&m_mid_cc}, m_rtol, m_atol); + + // Map back to the B-staggering (cell i -> node i). Only the top node + // plane in the nodal dimension reads a guard cell: the periodic + // image filled by FillBoundary here. + m_sol_cc.FillBoundary(m_geom.periodicity()); + +#ifdef AMREX_USE_OMP +#pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) +#endif + for (MFIter mfi(*x_mfarrvec[lev][comp], TilingIfNotGPU()); mfi.isValid(); ++mfi) + { + const Box& tbx = mfi.tilebox(); + Array4 const& xs = x_mfarrvec[lev][comp]->array(mfi); + Array4 const& sc = m_sol_cc.const_array(mfi); + ParallelFor(tbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + xs(i,j,k) = sc(i,j,k); + }); + } + } +} + +#endif diff --git a/Source/NonlinearSolvers/PreconditionerLibrary.H b/Source/NonlinearSolvers/PreconditionerLibrary.H index 6de960ce359..43dc3d00616 100644 --- a/Source/NonlinearSolvers/PreconditionerLibrary.H +++ b/Source/NonlinearSolvers/PreconditionerLibrary.H @@ -8,6 +8,7 @@ */ AMREX_ENUM(PreconditionerType, pc_curl_curl_mlmg, + pc_darwin_mlmg, pc_jacobi, pc_petsc, none