Output the potential phi on the macroparticles (#4599)

* Output the potential phi on the macroparticles

* Fix compilation

* Cleaner code by using `getParticlePosition`

* Add include statement

* Update const-ness

* Update code so that `phi` can be calculated

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Better boundary error message

* Define function to define tiles, for a NamedComponentParticleContainer

* Call parent class in WarpXParticleContainer

* Call `defineAllParticleTiles` in diagnostics

* Fix compilation issue

* Fix compilation for fixed precision

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add automated test

* Fix analysis script

* Update test so that they work correctly in RZ

* Update automated test

* Update implementation of `phi` on the particles

* Update test so as to trigger output of `phi`

* Update documentation

* Update Source/Diagnostics/ParticleDiag/ParticleDiag.H

* Apply suggestions from code review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

Docs/source/usage/parameters.rst

@@ -2676,7 +2676,8 @@ In-situ capabilities can be used by turning on Sensei or Ascent (provided they a
 * ``<diag_name>.<species_name>.variables`` (list of `strings` separated by spaces, optional)
     List of particle quantities to write to output.
     Choices are ``w`` for the particle weight and ``ux`` ``uy`` ``uz`` for the particle momenta.
-    By default, all particle quantities are written.
+    When using the lab-frame electrostatic solver, ``phi`` (electrostatic potential, on the macroparticles) is also available.
+    By default, all particle quantities (except ``phi``) are written.
     If ``<diag_name>.<species_name>.variables = none``, no particle data are written, except for particle positions, which are always included.
 
 * ``<diag_name>.<species_name>.random_fraction`` (`float`) optional

Examples/Tests/electrostatic_sphere/analysis_electrostatic_sphere.py

@@ -22,6 +22,8 @@
 
 import numpy as np
 import yt
+from openpmd_viewer import OpenPMDTimeSeries
+from scipy.constants import c
 from scipy.optimize import fsolve
 
 sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
@@ -125,7 +127,6 @@ def calculate_error(E_axis, xmin, dx, nx):
 
     return L2_error
 
-
 L2_error_x = calculate_error(Ex_axis, xmin, dx, nx)
 L2_error_y = calculate_error(Ey_axis, ymin, dy, ny)
 L2_error_z = calculate_error(Ez_axis, zmin, dz, nz)
@@ -137,6 +138,21 @@ def calculate_error(E_axis, xmin, dx, nx):
 assert L2_error_y < 0.05
 assert L2_error_z < 0.05
 
+# Check conservation of energy
+def return_energies(iteration):
+    ux, uy, uz, phi, m, q, w = ts.get_particle(['ux', 'uy', 'uz', 'phi', 'mass', 'charge', 'w'], iteration=iteration)
+    E_kinetic = (w*m*c**2 * (np.sqrt(1 + ux**2 + uy**2 + uz**2) - 1)).sum()
+    E_potential = 0.5*(w*q*phi).sum() # potential energy of particles in their own space-charge field: includes factor 1/2
+    return E_kinetic, E_potential
+ts = OpenPMDTimeSeries('./diags/diag2')
+if 'phi' in ts.avail_record_components['electron']:
+    # phi is only available when this script is run with the labframe poisson solver
+    print('Checking conservation of energy')
+    Ek_i, Ep_i = return_energies(0)
+    Ek_f, Ep_f = return_energies(30)
+    assert Ep_f < 0.7*Ep_i # Check that potential energy changes significantly
+    assert abs( (Ek_i + Ep_i) - (Ek_f + Ep_f) ) < 0.003 * (Ek_i + Ep_i) # Check conservation of energy
+
 # Checksum regression analysis
 test_name = os.path.split(os.getcwd())[1]
 checksumAPI.evaluate_checksum(test_name, filename)

Examples/Tests/electrostatic_sphere/inputs_3d

@@ -29,7 +29,13 @@ electron.profile = parse_density_function
 electron.density_function(x,y,z) = "(x*x + y*y + z*z < R0*R0)*n0"
 electron.momentum_distribution_type = at_rest
 
-diagnostics.diags_names = diag1
+diagnostics.diags_names = diag1 diag2
+
 diag1.intervals = 30
 diag1.diag_type = Full
 diag1.fields_to_plot = Ex Ey Ez rho
+
+diag2.intervals = 30
+diag2.diag_type = Full
+diag2.fields_to_plot = none
+diag2.format = openpmd

Examples/Tests/electrostatic_sphere/inputs_rz

@@ -30,7 +30,14 @@ electron.profile = parse_density_function
 electron.density_function(x,y,z) = "(x*x + y*y + z*z < R0*R0)*n0"
 electron.momentum_distribution_type = at_rest
 
-diagnostics.diags_names = diag1
+diagnostics.diags_names = diag1 diag2
+
 diag1.intervals = 30
 diag1.diag_type = Full
 diag1.fields_to_plot = Er Et Ez rho
+
+diag2.intervals = 30
+diag2.diag_type = Full
+diag2.fields_to_plot = none
+diag2.electron.variables = ux uy uz w phi
+diag2.format = openpmd

Regression/WarpX-tests.ini

@@ -554,7 +554,7 @@ analysisRoutine = Examples/Tests/electrostatic_sphere_eb/analysis_rz.py
 [ElectrostaticSphereLabFrame]
 buildDir = .
 inputFile = Examples/Tests/electrostatic_sphere/inputs_3d
-runtime_params = warpx.do_electrostatic=labframe
+runtime_params = warpx.do_electrostatic=labframe diag2.electron.variables=ux uy uz w phi
 dim = 3
 addToCompileString =
 cmakeSetupOpts = -DWarpX_DIMS=3

Source/Diagnostics/ParticleDiag/ParticleDiag.H

@@ -23,6 +23,7 @@ public:
     [[nodiscard]] PinnedMemoryParticleContainer* getPinnedParticleContainer() const { return m_pinned_pc; }
     [[nodiscard]] std::string getSpeciesName() const { return m_name; }
     amrex::Vector<int> m_plot_flags;
+    bool m_plot_phi = false; // Whether to output the potential phi on the particles
 
     bool m_do_random_filter  = false;
     bool m_do_uniform_filter = false;

Source/Diagnostics/ParticleDiag/ParticleDiag.cpp

@@ -35,12 +35,19 @@ ParticleDiag::ParticleDiag(
         if (variables[0] != "none"){
             const std::map<std::string, int> existing_variable_names = pc->getParticleComps();
             for (const auto& var : variables){
-                const auto search = existing_variable_names.find(var);
-                WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
-                    search != existing_variable_names.end(),
-                    "variables argument '" + var
-                    +"' is not an existing attribute for this species");
-                m_plot_flags[existing_variable_names.at(var)] = 1;
+                if (var == "phi") {
+                    // User requests phi on particle. This is *not* part of the variables that
+                    // the particle container carries, and is only added to particles during output.
+                    // Therefore, this case needs to be treated specifically.
+                    m_plot_phi = true;
+                } else {
+                    const auto search = existing_variable_names.find(var);
+                    WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
+                        search != existing_variable_names.end(),
+                        "variables argument '" + var
+                        +"' is not an existing attribute for this species");
+                    m_plot_flags[existing_variable_names.at(var)] = 1;
+                }
             }
         }
     }

Source/Diagnostics/WarpXOpenPMD.cpp

@@ -548,48 +548,6 @@ for (unsigned i = 0, n = particle_diags.size(); i < n; ++i) {
         pinned_pc->make_alike<amrex::PinnedArenaAllocator>() :
         pc->make_alike<amrex::PinnedArenaAllocator>();
 
-    // names of amrex::Real and int particle attributes in SoA data
-    amrex::Vector<std::string> real_names;
-    amrex::Vector<std::string> int_names;
-    amrex::Vector<int> int_flags;
-    amrex::Vector<int> real_flags;
-
-    // see openPMD ED-PIC extension for namings
-    // note: an underscore separates the record name from its component
-    //       for non-scalar records
-    // note: in RZ, we reconstruct x,y,z positions from r,z,theta in WarpX
-#if !defined (WARPX_DIM_1D_Z)
-    real_names.push_back("position_x");
-#endif
-#if defined (WARPX_DIM_3D) || defined(WARPX_DIM_RZ)
-    real_names.push_back("position_y");
-#endif
-    real_names.push_back("position_z");
-    real_names.push_back("weighting");
-    real_names.push_back("momentum_x");
-    real_names.push_back("momentum_y");
-    real_names.push_back("momentum_z");
-
-    // get the names of the real comps
-    real_names.resize(tmp.NumRealComps());
-    auto runtime_rnames = tmp.getParticleRuntimeComps();
-    for (auto const& x : runtime_rnames)
-    {
-        real_names[x.second+PIdx::nattribs] = detail::snakeToCamel(x.first);
-    }
-    // plot any "extra" fields by default
-    real_flags = particle_diags[i].m_plot_flags;
-    real_flags.resize(tmp.NumRealComps(), 1);
-    // and the names
-    int_names.resize(tmp.NumIntComps());
-    auto runtime_inames = tmp.getParticleRuntimeiComps();
-    for (auto const& x : runtime_inames)
-    {
-        int_names[x.second+0] = detail::snakeToCamel(x.first);
-    }
-    // plot by default
-    int_flags.resize(tmp.NumIntComps(), 1);
-
     const auto mass = pc->AmIA<PhysicalSpecies::photon>() ? PhysConst::m_e : pc->getMass();
     RandomFilter const random_filter(particle_diags[i].m_do_random_filter,
                                      particle_diags[i].m_random_fraction);
@@ -631,6 +589,76 @@ for (unsigned i = 0, n = particle_diags.size(); i < n; ++i) {
         particlesConvertUnits(ConvertDirection::SI_to_WarpX, pc, mass);
     }
 
+    // Gather the electrostatic potential (phi) on the macroparticles
+    if ( particle_diags[i].m_plot_phi ) {
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
+            (WarpX::electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrame) ||
+            (WarpX::electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrameElectroMagnetostatic),
+            "Output of the electrostatic potential (phi) on the particles was requested, "
+            "but this is only available for `warpx.do_electrostatic=labframe` or `labframe-electromagnetostatic`.");
+        // Using pinned PC indicates that the particles are not written at the same physical time (i.e. PIC iteration)
+        // that they were collected. This happens for diagnostics that use buffering (e.g. BackTransformed, BoundaryScraping).
+        // Here `phi` is gathered at the iteration when particles are written (not collected) and is thus mismatched.
+        // To avoid confusion, we raise an error in this case.
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
+            use_pinned_pc == false,
+            "Output of the electrostatic potential (phi) on the particles was requested, "
+            "but this is only available with `diag_type = Full`.");
+        tmp.AddRealComp("phi");
+        int const phi_index = tmp.getParticleComps().at("phi");
+        auto& warpx = WarpX::GetInstance();
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
+#endif
+        for (int lev=0; lev<=warpx.finestLevel(); lev++) {
+            const amrex::Geometry& geom = warpx.Geom(lev);
+            auto plo = geom.ProbLoArray();
+            auto dxi = geom.InvCellSizeArray();
+            amrex::MultiFab const& phi = warpx.getField( FieldType::phi_fp, lev, 0 );
+            storeFieldOnParticles( tmp, plo, dxi, phi, phi_index, lev );
+        }
+    }
+
+    // names of amrex::Real and int particle attributes in SoA data
+    amrex::Vector<std::string> real_names;
+    amrex::Vector<std::string> int_names;
+    amrex::Vector<int> int_flags;
+    amrex::Vector<int> real_flags;
+    // see openPMD ED-PIC extension for namings
+    // note: an underscore separates the record name from its component
+    //       for non-scalar records
+    // note: in RZ, we reconstruct x,y,z positions from r,z,theta in WarpX
+#if !defined (WARPX_DIM_1D_Z)
+    real_names.push_back("position_x");
+#endif
+#if defined (WARPX_DIM_3D) || defined(WARPX_DIM_RZ)
+    real_names.push_back("position_y");
+#endif
+    real_names.push_back("position_z");
+    real_names.push_back("weighting");
+    real_names.push_back("momentum_x");
+    real_names.push_back("momentum_y");
+    real_names.push_back("momentum_z");
+    // get the names of the real comps
+    real_names.resize(tmp.NumRealComps());
+    auto runtime_rnames = tmp.getParticleRuntimeComps();
+    for (auto const& x : runtime_rnames)
+    {
+        real_names[x.second+PIdx::nattribs] = detail::snakeToCamel(x.first);
+    }
+    // plot any "extra" fields by default
+    real_flags = particle_diags[i].m_plot_flags;
+    real_flags.resize(tmp.NumRealComps(), 1);
+    // and the names
+    int_names.resize(tmp.NumIntComps());
+    auto runtime_inames = tmp.getParticleRuntimeiComps();
+    for (auto const& x : runtime_inames)
+    {
+        int_names[x.second+0] = detail::snakeToCamel(x.first);
+    }
+    // plot by default
+    int_flags.resize(tmp.NumIntComps(), 1);
+
     // real_names contains a list of all real particle attributes.
     // real_flags is 1 or 0, whether quantity is dumped or not.
     DumpToFile(&tmp,
@@ -640,9 +668,8 @@ for (unsigned i = 0, n = particle_diags.size(); i < n; ++i) {
         int_flags,
         real_names, int_names,
         pc->getCharge(), pc->getMass(),
-        isBTD, isLastBTDFlush
-    );
-}
+        isBTD, isLastBTDFlush);
+    }
 }
 
 void

Source/Particles/ParticleIO.H

@@ -8,6 +8,9 @@
 #define PARTICLEIO_H_
 
 #include "Particles/WarpXParticleContainer.H"
+#include "Particles/Pusher/GetAndSetPosition.H"
+
+#include <ablastr/particles/NodalFieldGather.H>
 
 #include <AMReX_AmrParticles.H>
 #include <AMReX_ParIter.H>
@@ -77,4 +80,47 @@ particlesConvertUnits (ConvertDirection convert_direction, T_ParticleContainer*
     }
 }
 
+/** Gathers the field from a MultiFab to the macroparticles
+ * and stores it into a runtime component of the particle container
+ *
+ * @tparam T_ParticleContainer a WarpX particle container or AmrParticleContainer
+ * @param tmp the particle container on which to store the gathered field
+ * @param plo the coordinates of the low bound of the box, along each dimension
+ * @param dxi the inverse of the cell size, along each dimension
+ * @param phi a multifab that contains the field to be gathered
+ * @param phi_index the index of the runtime component where to store the gathered field
+ * @param lev the MR level
+ */
+template< typename T_ParticleContainer >
+void
+storeFieldOnParticles ( T_ParticleContainer& tmp,
+    amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
+    amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi,
+    amrex::MultiFab const& phi,
+    int const phi_index, int const lev ) {
+
+    using PinnedParIter = typename T_ParticleContainer::ParIterType;
+
+    for (PinnedParIter pti(tmp, lev); pti.isValid(); ++pti) {
+
+        auto phi_grid = phi[pti].array();
+        const auto getPosition = GetParticlePosition<PIdx>(pti);
+        amrex::ParticleReal* phi_particle_arr = pti.GetStructOfArrays().GetRealData(phi_index).dataPtr();
+
+        // Loop over the particles and update their position
+        amrex::ParallelFor( pti.numParticles(),
+            [=] AMREX_GPU_DEVICE (long ip) {
+
+                amrex::ParticleReal xp, yp, zp;
+                getPosition(ip, xp, yp, zp);
+                int i, j, k;
+                amrex::Real W[AMREX_SPACEDIM][2];
+                ablastr::particles::compute_weights(xp, yp, zp, plo, dxi, i, j, k, W);
+                amrex::Real const phi_value  = ablastr::particles::interp_field_nodal(i, j, k, W, phi_grid);
+                phi_particle_arr[ip] = phi_value;
+            }
+        );
+    }
+}
+
 #endif /* PARTICLEIO_H_ */