Add theta-implicit hybrid-PIC solver

Docs/source/usage/parameters.rst

@@ -352,6 +352,29 @@ Overall simulation parameters
         - `Angus et al., An implicit particle code with exact energy and charge conservation for electromagnetic studies of dense plasmas <https://doi.org/10.1016/j.jcp.2023.112383>`__.
         - `Angus et al., An implicit particle code with exact energy and charge conservation for studies of dense plasmas in axisymmetric geometries <https://doi.org/10.1016/j.jcp.2024.113427>`__.
 
+    * ``theta_implicit_hybrid``: Use a :math:`\theta`-implicit solver for the hybrid-PIC (Ohm's law) model.
+
+      The fields and particles are advanced as a single coupled nonlinear system: at each nonlinear
+      iteration the particles are re-advanced with the iterate electric field, the ion current and
+      charge density are re-deposited, the magnetic field is updated through Faraday's law, and the
+      electric field is re-evaluated from the generalized Ohm's law. Since the whistler-wave CFL
+      condition of the explicit hybrid advance does not apply, no field sub-stepping is needed
+      (:pp:param:`hybrid_pic_model.substeps` is ignored).
+
+      - Requires :pp:param:`algo.maxwell_solver` = ``hybrid`` and one of the ``direct``, ``villasenor``
+        or ``esirkepov`` current depositions.
+      - The time-biasing parameter ``implicit_evolve.theta`` and the nonlinear solver parameters
+        (``implicit_evolve.nonlinear_solver``, ``picard.*``, ``newton.*``, ``gmres.*``,
+        ``implicit_evolve.max_particle_iterations``, ``implicit_evolve.particle_tolerance``,
+        ``implicit_evolve.particle_suborbits``) are shared with ``theta_implicit_em``; see above.
+      - The electron pressure is evaluated from the algebraic equation-of-state closure
+        (see :ref:`the hybrid-PIC theory section <theory-kinetic-fluid-hybrid-model>`).
+        ``hybrid_pic_model.solve_electron_energy_equation``, mass matrices
+        (``implicit_evolve.use_mass_matrices_jacobian``) and external fields from vector potentials
+        are not yet supported with this scheme.
+      - The method is described in `Kumar et al., A theta-implicit hybrid particle-in-cell algorithm
+        for kinetic simulations of magnetized plasmas <https://dx.doi.org/10.2139/ssrn.7346445>`__.
+
     * ``semi_implicit_em``: Use an approximately energy conserving semi-implicit electromagnetic solver.
 
       - Difference with ``theta_implicit_em`` is that light waves are treated explicit just as in the standard FDTD method. Consequently, this method has the CFL limitation :math:`c\Delta t < 1/\sqrt( \sum_i 1/\Delta x_i^2 )`.

Examples/Tests/magnetized_plasma_modes/CMakeLists.txt

@@ -11,6 +11,26 @@ add_warpx_test(
     OFF  # dependency
 )
 
+add_warpx_test(
+    test_1d_ohm_solver_em_modes_implicit_picmi  # name
+    1  # dims
+    2  # nprocs
+    "inputs_test_em_modes_picmi.py --test --dim 1 --bdir z --ohm --implicit"  # inputs
+    "analysis.py --analyze_ohm_sim"  # analysis
+    "analysis_default_regression.py --path diags/field_diag000050"  # checksum
+    OFF  # dependency
+)
+
+add_warpx_test(
+    test_1d_ohm_solver_em_modes_perp_implicit_picmi  # name
+    1  # dims
+    2  # nprocs
+    "inputs_test_em_modes_picmi.py --test --dim 1 --bdir x --ohm --implicit"  # inputs
+    "analysis.py --analyze_ohm_sim"  # analysis
+    "analysis_default_regression.py --path diags/field_diag000050"  # checksum
+    OFF  # dependency
+)
+
 add_warpx_test(
     test_1d_ohm_solver_em_modes_picmi  # name
     1  # dims

Examples/Tests/magnetized_plasma_modes/analysis.py

@@ -439,3 +439,20 @@ def get_analytic_L_mode(w):
     )
 if not sim.test:
     plt.show()
+
+if not is_darwin and getattr(sim, "use_implicit", False):
+    # Check the total (field + ion kinetic) energy drift of the theta-implicit
+    # hybrid solver. At theta = 0.5 the field-particle energy exchange of the
+    # scheme is discretely conservative; the remaining drift comes from the
+    # algebraic electron pressure closure (whose grad(Pe) work on the ions has
+    # no paired electron internal energy ledger) and from the small physical
+    # resistive dissipation, both of which are bounded and small here.
+    field_energy = np.loadtxt("diags/field_energy.txt", skiprows=1)
+    part_energy = np.loadtxt("diags/part_energy.txt", skiprows=1)
+    W_tot = field_energy[:, 2] + part_energy[:, 2]
+    max_drift = np.max(np.abs(W_tot - W_tot[0])) / W_tot[0]
+    print(f"Maximum relative energy drift: {max_drift:.3e}")
+    tolerance = 5e-4
+    assert max_drift < tolerance, (
+        f"Total energy drift {max_drift:.3e} exceeds the tolerance {tolerance:.1e}"
+    )

Examples/Tests/magnetized_plasma_modes/inputs_test_em_modes_picmi.py

@@ -16,6 +16,7 @@
 from mpi4py import MPI as mpi
 
 from pywarpx import callbacks, libwarpx, picmi
+from pywarpx import particles as pywarpx_particles
 
 constants = picmi.constants
 
@@ -95,6 +96,7 @@ def __init__(
         verbose,
         include_es_solver=False,
         use_rkf45=False,
+        use_implicit=False,
     ):
         """Get input parameters for the specific case desired."""
         self.solver = solver
@@ -104,6 +106,7 @@ def __init__(
         self.verbose = verbose or self.test
         self.include_es_solver = include_es_solver
         self.use_rkf45 = use_rkf45
+        self.use_implicit = use_implicit
 
         # sanity check
         assert dim > 0 and dim < 4, f"{dim}-dimensions not a valid input"
@@ -117,6 +120,18 @@ def __init__(
         # the initial B-field and the dimensionality
         self.get_simulation_parameters()
 
+        if self.test and self.use_implicit:
+            # Every GMRES iteration of the (unpreconditioned) implicit solver
+            # evaluates the full nonlinear residual, so the CI case is kept
+            # small: fewer/coarser cells and particles than the explicit CI
+            # case, and a time step ~3x beyond the whistler CFL limit
+            # dt < dz^2/(pi*l_i*v_A) of the explicit advance, taken in a
+            # single implicit step with no field subcycling.
+            self.Nz = 64
+            self.NPPC = 64
+            self.DZ = 0.25
+            self.DT = 0.01
+
         # calculate various plasma parameters based on the simulation input
         self.get_plasma_quantities()
 
@@ -139,7 +154,12 @@ def __init__(
                 self.total_steps = int(self.LT / self.DT)
         else:
             # if this is a test case run for only a small number of steps
-            self.total_steps = 50 if self.solver == "darwin" else 250
+            if self.solver == "darwin":
+                self.total_steps = 50
+            elif self.use_implicit:
+                self.total_steps = 50
+            else:
+                self.total_steps = 250
 
         if self.solver == "darwin":
             self.diag_steps = 3
@@ -331,6 +351,27 @@ def setup_run(self):
             )
             simulation.solver = self.solver_obj
 
+            if self.use_implicit:
+                simulation.evolve_scheme = picmi.ThetaImplicitHybridEvolveScheme(
+                    nonlinear_solver=picmi.NewtonNonlinearSolver(
+                        verbose=self.verbose,
+                        relative_tolerance=1e-6,
+                        max_iterations=20,
+                        require_convergence=True,
+                        linear_solver=picmi.GMRESLinearSolver(
+                            verbose_int=0,
+                            relative_tolerance=1e-6,
+                            max_iterations=1000,
+                            restart_length=100,
+                        ),
+                    ),
+                    theta=0.5,
+                )
+                # Particles may straddle a grid-box face at the time-centered
+                # deposition inside the nonlinear solve; allow for one grid
+                # crossing so the rho/J guard cells cover their deposition.
+                pywarpx_particles.max_grid_crossings = 2
+
         B_ext = picmi.AnalyticInitialField(
             Bx_expression=self.Bx, By_expression=self.By, Bz_expression=self.Bz
         )
@@ -429,7 +470,7 @@ def setup_run(self):
                     "[3]Ez_lev0-(V/m) [4]Bx_lev0-(T) [5]By_lev0-(T)\n"
                 )
 
-        if self.solver == "darwin":
+        if self.solver == "darwin" or self.use_implicit:
             write_dir = "diags/"
             field_energy = picmi.ReducedDiagnostic(
                 diag_type="FieldEnergy",
@@ -541,6 +582,11 @@ def _record_average_fields(self):
     help="Ohm only: use adaptive RKF45 subcycling for the B-field update",
     action="store_true",
 )
+parser.add_argument(
+    "--implicit",
+    help="Ohm only: use the theta-implicit hybrid evolve scheme",
+    action="store_true",
+)
 parser.add_argument(
     "-v",
     "--verbose",
@@ -558,5 +604,6 @@ def _record_average_fields(self):
     verbose=args.verbose,
     include_es_solver=args.include_es_solver,
     use_rkf45=args.use_rkf45,
+    use_implicit=args.implicit,
 )
 simulation.step()

Python/pywarpx/picmi.py

@@ -2078,6 +2078,34 @@ def solver_scheme_initialize_inputs(self):
         self.nonlinear_solver.nonlinear_solver_initialize_inputs()
 
 
+class ThetaImplicitHybridEvolveScheme(picmistandard.base._ClassWithInit):
+    """
+    Sets up the "theta implicit" hybrid-PIC evolve scheme. This requires the
+    hybrid-PIC (Ohm's law) field solver.
+
+    Parameters
+    ----------
+    nonlinear_solver: nonlinear solver instance
+        The nonlinear solver to use for the iterations
+
+    theta: float, optional
+        The "theta" parameter, determining the level of implicitness
+    """
+
+    def __init__(self, nonlinear_solver, theta=None):
+        self.nonlinear_solver = nonlinear_solver
+        self.theta = theta
+
+        assert isinstance(nonlinear_solver, NonlinearSolverBase)
+
+    def solver_scheme_initialize_inputs(self):
+        pywarpx.algo.evolve_scheme = "theta_implicit_hybrid"
+        implicit_evolve = pywarpx.warpx.get_bucket("implicit_evolve")
+        implicit_evolve.theta = self.theta
+
+        self.nonlinear_solver.nonlinear_solver_initialize_inputs()
+
+
 class SemiImplicitEMEvolveScheme(picmistandard.base._ClassWithInit):
     """
     Sets up the "semi-implicit" electromagnetic evolve scheme

Regression/Checksum/benchmarks_json/test_1d_ohm_solver_em_modes_implicit_picmi.json

@@ -0,0 +1,20 @@
+{
+  "ions": {
+    "particle_momentum_x": 6.326138045981961e-22,
+    "particle_momentum_y": 6.307852461681281e-22,
+    "particle_momentum_z": 6.151869795493248e-22,
+    "particle_position_x": 2.2420806128986115,
+    "particle_weight": 6.594142739277304e+20
+  },
+  "lev=0": {
+    "Bx": 0.026766665175599412,
+    "By": 0.039440176598767196,
+    "Bz": 16.0,
+    "Ex": 2051.0444731149764,
+    "Ey": 2193.5557387135805,
+    "Ez": 1960.5981881554917,
+    "jx_displacement": 320593775.8826201,
+    "jy_displacement": 372833136.8308284,
+    "jz_displacement": 760627461.4128174
+  }
+}

Regression/Checksum/benchmarks_json/test_1d_ohm_solver_em_modes_perp_implicit_picmi.json

@@ -0,0 +1,20 @@
+{
+  "ions": {
+    "particle_momentum_x": 8.016845243683861e-20,
+    "particle_momentum_y": 8.024572617022076e-20,
+    "particle_momentum_z": 7.792795013905813e-20,
+    "particle_position_x": 89.39859345506427,
+    "particle_weight": 6.618763146287724e+19
+  },
+  "lev=0": {
+    "Bx": 16.0,
+    "By": 0.0,
+    "Bz": 0.0,
+    "Ex": 0.0,
+    "Ey": 66292.97288979663,
+    "Ez": 323088.32409385114,
+    "jx_displacement": 61712446.561408125,
+    "jy_displacement": 215235629.30999476,
+    "jz_displacement": 70263879.4200543
+  }
+}

Source/Evolve/WarpXEvolve.cpp

@@ -326,8 +326,11 @@ WarpX::Evolve (int numsteps)
             ExecutePythonCallback("afterEsolve");
         }
 
-        // Hybrid-PIC case
-        if (electromagnetic_solver_id == ElectromagneticSolverAlgo::HybridPIC) {
+        // Hybrid-PIC case. With the theta-implicit hybrid evolve scheme the
+        // fields are already advanced self-consistently with the particles in
+        // the implicit solver's OneStep, so this explicit field update is skipped.
+        if (electromagnetic_solver_id == ElectromagneticSolverAlgo::HybridPIC &&
+            !m_implicit_solver) {
             ExecutePythonCallback("beforeEsolve");
             // The particles are now at p^{n+1/2} and x^{n+1}. The fields
             // are updated according to the hybrid-PIC scheme (Ohm's law

Source/FieldSolver/FiniteDifferenceSolver/FiniteDifferenceSolver.H

@@ -156,6 +156,9 @@ class FiniteDifferenceSolver
           * \param[in] lev  level number for the calculation
           * \param[in] hybrid_model instance of the hybrid-PIC model
           * \param[in] solve_for_Faraday boolean flag for whether the E-field is solved to be used in Faraday's equation
+          * \param[in] solve_for_implicit boolean flag to include the electron pressure
+          *            gradient term even when solve_for_Faraday is true (used by the
+          *            implicit solver, where grad(Pe) is retained in the Newton residual)
           */
         void HybridPICSolveE ( ablastr::fields::VectorField const& Efield,
                                ablastr::fields::VectorField & Jfield,
@@ -165,7 +168,8 @@ class FiniteDifferenceSolver
                                amrex::MultiFab const& Pefield,
                                std::array< std::unique_ptr<amrex::iMultiFab>,3> const& eb_update_E,
                                int lev, HybridPICModel const* hybrid_model,
-                               bool solve_for_Faraday );
+                               bool solve_for_Faraday,
+                               bool solve_for_implicit = false );
 
         /**
           * \brief Calculation of total current using Ampere's law (without
@@ -346,7 +350,8 @@ class FiniteDifferenceSolver
             amrex::MultiFab const& Pefield,
             std::array< std::unique_ptr<amrex::iMultiFab>,3> const& eb_update_E,
             int lev, HybridPICModel const* hybrid_model,
-            bool solve_for_Faraday );
+            bool solve_for_Faraday,
+            bool solve_for_implicit = false );
 
         template<typename T_Algo>
         void CalculateCurrentAmpereCylindrical (
@@ -444,7 +449,8 @@ class FiniteDifferenceSolver
             amrex::MultiFab const& rhofield,
             amrex::MultiFab const& Pefield,
             int lev, HybridPICModel const* hybrid_model,
-            bool solve_for_Faraday );
+            bool solve_for_Faraday,
+            bool solve_for_implicit = false );
 
         template<typename T_Algo>
         void CalculateCurrentAmpereSpherical (
@@ -595,7 +601,8 @@ class FiniteDifferenceSolver
             amrex::MultiFab const& Pefield,
             std::array< std::unique_ptr<amrex::iMultiFab>,3> const& eb_update_E,
             int lev, HybridPICModel const* hybrid_model,
-            bool solve_for_Faraday );
+            bool solve_for_Faraday,
+            bool solve_for_implicit = false );
 
         template<typename T_Algo>
         void CalculateCurrentAmpereCartesian (

Source/FieldSolver/FiniteDifferenceSolver/HybridPICModel/HybridPICModel.H

@@ -119,23 +119,26 @@ public:
         ablastr::fields::MultiLevelVectorField const& Bfield,
         ablastr::fields::MultiLevelScalarField const& rhofield,
         amrex::Vector<std::array< std::unique_ptr<amrex::iMultiFab>,3 > >& eb_update_E,
-        bool solve_for_Faraday) const;
+        bool solve_for_Faraday,
+        bool solve_for_implicit = false) const;
 
     void HybridPICSolveE (
         ablastr::fields::VectorField const& Efield,
         ablastr::fields::VectorField const& Jfield,
         ablastr::fields::VectorField const& Bfield,
         amrex::MultiFab const& rhofield,
         std::array< std::unique_ptr<amrex::iMultiFab>,3 >& eb_update_E,
-        int lev, bool solve_for_Faraday) const;
+        int lev, bool solve_for_Faraday,
+        bool solve_for_implicit = false) const;
 
     void HybridPICSolveE (
         ablastr::fields::VectorField const& Efield,
         ablastr::fields::VectorField const& Jfield,
         ablastr::fields::VectorField const& Bfield,
         amrex::MultiFab const& rhofield,
         std::array< std::unique_ptr<amrex::iMultiFab>,3 >& eb_update_E,
-        int lev, PatchType patch_type, bool solve_for_Faraday) const;
+        int lev, PatchType patch_type, bool solve_for_Faraday,
+        bool solve_for_implicit = false) const;
 
     void BfieldEvolve (
         ablastr::fields::MultiLevelVectorField const& Bfield,