Add a factored-Laplacian MLMG preconditioner for the Darwin GMRES solve

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:
 

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

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()

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

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
   }
 }