Add strang implicit spectral em redo (#5027)

This implements use of the PSATD field advance coupled with the implicit
solver, using a Strang split advance.

- Advect Maxwell using PSATD with no J, ½ step
- Advance particles along with dE/dt = -J implicitly, full step,
iterating
- Advect Maxwell using PSATD with no J, ½ step

This requires the input parameter psatd.update_with_rho = 0.
With psatd.periodic_single_box_fft = 1, exact energy conservation is
obtained. Otherwise good conservation is seen, but not exact (will
depend on parameters).
Convergence is found for wpedt <= 1.9 (compared to wpedt < 0.25 for
FDTD).

This PR replaces PR #4662.

A task for a future PR would be to implement specialized source free
spectral advance routines (as noted in source comments).

Docs/source/usage/parameters.rst

@@ -93,6 +93,14 @@ Overall simulation parameters
       The PS-JFNK method is described in `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>`__ . (The version implemented in WarpX is an updated version that includes the relativistic gamma factor for the particles.) Also see `Chen et al., An energy- and charge-conserving, implicit, electrostatic particle-in-cell algorithm. <https://doi.org/10.1016/j.jcp.2011.05.031>`__ .
       Exact energy conservation requires that the interpolation stencil used for the field gather match that used for the current deposition. ``algo.current_deposition = direct`` must be used with ``interpolation.galerkin_scheme = 0``, and ``algo.current_deposition = Esirkepov`` must be used with ``interpolation.galerkin_scheme = 1``. If using ``algo.current_deposition = villasenor``, the corresponding field gather routine will automatically be selected and the ``interpolation.galerkin_scheme`` flag does not need to be specified. The Esirkepov and villasenor deposition schemes are charge-conserving.
 
+    * ``strang_implicit_spectral_em``: Use a fully implicit electromagnetic solver. All of the comments for ``theta_implicit_em``
+      above apply here as well (except that theta is fixed to 0.5 and that charge will not be conserved).
+      In this version, the advance is Strang split, with a half advance of the source free Maxwell's equation (with a spectral solver), a full advance of the particles plus longitudinal E field, and a second half advance of the source free Maxwell's equations.
+      The advantage of this method is that with the Spectral advance of the fields, it is dispersionless.
+      Note that exact energy convergence is achieved only with one grid block and ``psatd.periodic_single_box_fft == 1``. Otherwise,
+      the energy convservation is spoiled because of the inconsistency of the periodic assumption of the spectral solver and the
+      non-periodic behavior of the individual blocks.
+
     * ``semi_implicit_em``: Use an approximately energy conserving semi-implicit electromagnetic solver. Choices for the nonlinear solver include a Picard iteration scheme and particle-suppressed JFNK.
       Note that this method has the CFL limitation :math:`\Delta t < c/\sqrt( \sum_i 1/\Delta x_i^2 )`. The Picard solver for this method can only be expected to work well when :math:`\omega_{pe} \Delta t` is less than one.
       The method is described in `Chen et al., A semi-implicit, energy- and charge-conserving particle-in-cell algorithm for the relativistic Vlasov-Maxwell equations <https://doi.org/10.1016/j.jcp.2020.109228>`__.
@@ -105,16 +113,16 @@ Overall simulation parameters
     exactly energy conserving, but the solver may perform better.
 
 * ``implicit_evolve.nonlinear_solver`` (`string`, default: None)
-    When `algo.evolve_scheme` is either `theta_implicit_em` or `semi_implicit_em`, this sets the nonlinear solver used
+    When `algo.evolve_scheme` is either `theta_implicit_em`, `strang_implicit_spectral_em`, or `semi_implicit_em`, this sets the nonlinear solver used
     to advance the field-particle system in time. Options are `picard` or `newton`.
 
 * ``implicit_evolve.max_particle_iterations`` (`integer`, default: 21)
-    When `algo.evolve_scheme` is either `theta_implicit_em` or `semi_implicit_em` and `implicit_evolve.nonlinear_solver = newton`
+    When `algo.evolve_scheme` is either `theta_implicit_em`, `strang_implicit_spectral_em`, or `semi_implicit_em` and `implicit_evolve.nonlinear_solver = newton`
     , this sets the maximum number of iterations for the method used to obtain a self-consistent update of the particles at
     each iteration in the JFNK process.
 
 * ``implicit_evolve.particle_tolerance`` (`float`, default: 1.e-10)
-    When `algo.evolve_scheme` is either `theta_implicit_em` or `semi_implicit_em` and `implicit_evolve.nonlinear_solver = newton`
+    When `algo.evolve_scheme` is either `theta_implicit_em`, `strang_implicit_spectral_em`, or `semi_implicit_em` and `implicit_evolve.nonlinear_solver = newton`
     , this sets the relative tolerance for the iterative method used to obtain a self-consistent update of the particles at
     each iteration in the JFNK process.
 

Examples/Tests/implicit/CMakeLists.txt

@@ -50,3 +50,15 @@ add_warpx_test(
     diags/diag1000020  # output
     OFF  # dependency
 )
+
+if(WarpX_FFT)
+    add_warpx_test(
+        test_2d_theta_implicit_strang_psatd  # name
+        2  # dims
+        2  # nprocs
+        inputs_test_2d_theta_implicit_strang_psatd  # inputs
+        analysis_2d_psatd.py  # analysis
+        diags/diag1000020  # output
+        OFF  # dependency
+    )
+endif()

Examples/Tests/implicit/analysis_2d_psatd.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+
+# Copyright 2024 Justin Angus, David Grote
+#
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+#
+# This is a script that analyses the simulation results from the script `inputs_vandb_2d`.
+# This simulates a 2D periodic plasma using the implicit solver
+# with the Villasenor deposition using shape factor 2.
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(1, "../../../../warpx/Regression/Checksum/")
+import checksumAPI
+
+# this will be the name of the plot file
+fn = sys.argv[1]
+
+field_energy = np.loadtxt("diags/reducedfiles/field_energy.txt", skiprows=1)
+particle_energy = np.loadtxt("diags/reducedfiles/particle_energy.txt", skiprows=1)
+
+total_energy = field_energy[:, 2] + particle_energy[:, 2]
+
+delta_E = (total_energy - total_energy[0]) / total_energy[0]
+max_delta_E = np.abs(delta_E).max()
+
+# This case should have near machine precision conservation of energy
+tolerance_rel_energy = 2.1e-14
+
+print(f"max change in energy: {max_delta_E}")
+print(f"tolerance: {tolerance_rel_energy}")
+
+assert max_delta_E < tolerance_rel_energy
+
+test_name = os.path.split(os.getcwd())[1]
+checksumAPI.evaluate_checksum(test_name, fn)

Examples/Tests/implicit/inputs_test_2d_theta_implicit_strang_psatd

@@ -0,0 +1,98 @@
+#################################
+########## CONSTANTS ############
+#################################
+
+my_constants.n0 = 1.e30          # m^-3
+my_constants.nz = 40
+my_constants.Ti = 100.           # eV
+my_constants.Te = 100.           # eV
+my_constants.wpe = q_e*sqrt(n0/(m_e*epsilon0))
+my_constants.de0 = clight/wpe
+my_constants.nppcz = 10          # number of particles/cell in z
+my_constants.dt = 0.1/wpe        # s
+
+#################################
+####### GENERAL PARAMETERS ######
+#################################
+max_step = 20
+amr.n_cell = nz nz
+amr.max_grid_size = nz
+amr.max_level = 0
+geometry.dims = 2
+geometry.prob_lo = 0.0  0.0  # physical domain
+geometry.prob_hi = 10.0*de0  10.0*de0
+
+#################################
+####### Boundary condition ######
+#################################
+boundary.field_lo = periodic periodic
+boundary.field_hi = periodic periodic
+
+#################################
+############ NUMERICS ###########
+#################################
+warpx.serialize_initial_conditions = 1
+warpx.verbose = 1
+warpx.const_dt = dt
+#warpx.cfl = 0.5656
+warpx.use_filter = 0
+
+algo.maxwell_solver = psatd
+algo.evolve_scheme = strang_implicit_spectral_em
+implicit_evolve.nonlinear_solver = "picard"
+
+picard.verbose = true
+picard.max_iterations = 9
+picard.relative_tolerance = 0.0
+picard.absolute_tolerance = 0.0
+picard.require_convergence = false
+
+algo.particle_pusher = "boris"
+
+algo.particle_shape = 2
+algo.current_deposition = direct
+algo.charge_deposition = standard
+algo.field_gathering = energy-conserving
+interpolation.galerkin_scheme = 0
+
+psatd.periodic_single_box_fft = 1
+psatd.update_with_rho = 0
+
+#################################
+############ PLASMA #############
+#################################
+particles.species_names = electrons protons
+
+electrons.species_type = electron
+electrons.injection_style = "NUniformPerCell"
+electrons.num_particles_per_cell_each_dim = nppcz nppcz
+electrons.profile = constant
+electrons.density = n0
+electrons.momentum_distribution_type = gaussian
+electrons.ux_th = sqrt(Te*q_e/m_e)/clight
+electrons.uy_th = sqrt(Te*q_e/m_e)/clight
+electrons.uz_th = sqrt(Te*q_e/m_e)/clight
+
+protons.species_type = proton
+protons.injection_style = "NUniformPerCell"
+protons.num_particles_per_cell_each_dim = nppcz nppcz
+protons.profile = constant
+protons.density = n0
+protons.momentum_distribution_type = gaussian
+protons.ux_th = sqrt(Ti*q_e/m_p)/clight
+protons.uy_th = sqrt(Ti*q_e/m_p)/clight
+protons.uz_th = sqrt(Ti*q_e/m_p)/clight
+
+# Diagnostics
+diagnostics.diags_names = diag1
+diag1.intervals = 20
+diag1.diag_type = Full
+diag1.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz rho divE
+diag1.electrons.variables = x z w ux uy uz
+diag1.protons.variables = x z w ux uy uz
+
+warpx.reduced_diags_names = particle_energy field_energy
+particle_energy.type = ParticleEnergy
+particle_energy.intervals = 1
+field_energy.type = FieldEnergy
+field_energy.intervals = 1

Regression/Checksum/benchmarks_json/test_2d_theta_implicit_strang_psatd.json

@@ -0,0 +1,31 @@
+{
+  "lev=0": {
+    "Bx": 60642.062637340816,
+    "By": 89855.09371265332,
+    "Bz": 54561.47120738846,
+    "Ex": 81536346169528.28,
+    "Ey": 13888711042388.54,
+    "Ez": 86853122458391.0,
+    "divE": 9.492653438830812e+22,
+    "jx": 2.5941826848709296e+19,
+    "jy": 2.9929071160915993e+19,
+    "jz": 2.692985701872205e+19,
+    "rho": 851978517887.51
+  },
+  "electrons": {
+    "particle_momentum_x": 4.864385990952573e-19,
+    "particle_momentum_y": 4.879723483907468e-19,
+    "particle_momentum_z": 4.865564630727981e-19,
+    "particle_position_x": 0.004250851253052539,
+    "particle_position_y": 0.0042513622554793,
+    "particle_weight": 2823958719279159.5
+  },
+  "protons": {
+    "particle_momentum_x": 2.0934469726422704e-17,
+    "particle_momentum_y": 2.0929630794865952e-17,
+    "particle_momentum_z": 2.093085625201003e-17,
+    "particle_position_x": 0.004251276208274589,
+    "particle_position_y": 0.004251274670600805,
+    "particle_weight": 2823958719279159.5
+  }
+}

Source/FieldSolver/ImplicitSolvers/CMakeLists.txt

@@ -5,6 +5,7 @@ foreach(D IN LISTS WarpX_DIMS)
         ImplicitSolver.cpp
         SemiImplicitEM.cpp
         ThetaImplicitEM.cpp
+        StrangImplicitSpectralEM.cpp
         WarpXImplicitOps.cpp
         WarpXSolverVec.cpp
     )

Source/FieldSolver/ImplicitSolvers/ImplicitSolverLibrary.H

@@ -9,5 +9,6 @@
 
 #include "SemiImplicitEM.H"  // IWYU pragma: export
 #include "ThetaImplicitEM.H"  // IWYU pragma: export
+#include "StrangImplicitSpectralEM.H"  // IWYU pragma: export
 
 #endif

Source/FieldSolver/ImplicitSolvers/Make.package

@@ -1,6 +1,7 @@
 CEXE_sources += ImplicitSolver.cpp
 CEXE_sources += SemiImplicitEM.cpp
 CEXE_sources += ThetaImplicitEM.cpp
+CEXE_sources += StrangImplicitSpectralEM.cpp
 CEXE_sources += WarpXImplicitOps.cpp
 CEXE_sources += WarpXSolverVec.cpp
 

Source/FieldSolver/ImplicitSolvers/StrangImplicitSpectralEM.H

@@ -0,0 +1,107 @@
+/* Copyright 2024 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef STRANG_IMPLICIT_SPECTRALEM_H_
+#define STRANG_IMPLICIT_SPECTRALEM_H_
+
+#include "FieldSolver/ImplicitSolvers/WarpXSolverVec.H"
+
+#include <AMReX_Array.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_REAL.H>
+
+#include "ImplicitSolver.H"
+
+/** @file
+ *  Implicit spectral electromagnetic time solver class. This is a fully implicit
+ *  algorithm where both the fields and particles are treated implicitly.
+ *
+ *  The time stencil is
+ *    Advance (Eg^n, Bg^n) -> (Eg^{n+1/2}, Bg^{n+1/2}) source free // E transverse
+ *    Iterate:
+ *      Eg^{n+1} = Eg^n + c^2*dt*( - mu0*Jg^{n+1/2} ) // E longitudinal
+ *      xp^{n+1} = xp^n + dt*up^{n+1/2}/(0.5*(gammap^n + gammap^{n+1}))
+ *      up^{n+1} = up^n + dt*qp/mp*(Ep^{n+1/2} + up^{n+1/2}/gammap^{n+1/2} x Bp^{n+1/2})
+ *    Advance (Eg^n+1/2, Bg^n+1/2) -> (Eg^{n+1}, Bg^{n+1}) source free // E transverse
+ *
+ *  The algorithm is exactly energy conserving only with a single box, periodic fft (psatd.periodic_single_box_fft = 1).
+ *  With multiple boxes, energy is not conserved since the ffts in each box assumes periodic in the box which
+ *  is not consistent with the current.
+ *  The algorithm is numerially stable for any time step.
+ *  I.e., the CFL condition for light waves does not
+ *  have to be satisifed and the time step is not limited by the plasma period. However, how
+ *  efficiently the algorithm can use large time steps depends strongly on the nonlinear solver.
+ *  Furthermore, the time step should always be such that particles do not travel outside the
+ *  ghost region of the box they live in, which is an MPI-related limitation. The time step
+ *  is always limited by the need to resolve the appropriate physics.
+ *
+ */
+
+class StrangImplicitSpectralEM : public ImplicitSolver
+{
+public:
+
+    StrangImplicitSpectralEM() = default;
+
+    ~StrangImplicitSpectralEM() override = default;
+
+    // Prohibit Move and Copy operations
+    StrangImplicitSpectralEM(const StrangImplicitSpectralEM&) = delete;
+    StrangImplicitSpectralEM& operator=(const StrangImplicitSpectralEM&) = delete;
+    StrangImplicitSpectralEM(StrangImplicitSpectralEM&&) = delete;
+    StrangImplicitSpectralEM& operator=(StrangImplicitSpectralEM&&) = delete;
+
+    void Define ( WarpX*  a_WarpX ) override;
+
+    void PrintParameters () const override;
+
+    void OneStep ( amrex::Real  a_time,
+                   amrex::Real  a_dt,
+                   int          a_step ) override;
+
+    void ComputeRHS ( WarpXSolverVec&  a_RHS,
+                const WarpXSolverVec&  a_E,
+                      amrex::Real      a_time,
+                      amrex::Real      a_dt,
+                      int              a_nl_iter,
+                      bool             a_from_jacobian ) override;
+
+private:
+
+    /**
+     * \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
+     * WarpXSolverVec type is that it must have the same size and have the same
+     * centering of the data as the variable being solved for, which is E here.
+     * For example, if using a Yee grid then a container for curlB could be a
+     * WarpXSovlerVec, but magnetic field B should not be.
+     */
+    WarpXSolverVec m_E, m_Eold;
+
+    /**
+     * \brief B is a derived variable from E. Need to save Bold to update B during
+     * the iterative nonlinear solve for E. Bold is owned here, but only used by WarpX.
+     * It is not used directly by the nonlinear solver, nor is it the same size as the
+     * solver vector (size E), and so it should not be WarpXSolverVec type.
+     */
+    amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>, 3 > > m_Bold;
+
+    /**
+     * \brief Update the E and B fields owned by WarpX
+     */
+    void UpdateWarpXFields ( WarpXSolverVec const&  a_E,
+                             amrex::Real            a_time,
+                             amrex::Real            a_dt );
+
+    /**
+     * \brief Nonlinear solver is for the time-centered values of E. After
+     * the solver, need to use m_E and m_Eold to compute E^{n+1}
+     */
+    void FinishFieldUpdate ( amrex::Real  a_new_time );
+
+};
+
+#endif