Improve rho diagnostic

Docs/source/visualization/advanced.rst

@@ -17,7 +17,7 @@ Dump additional data
 --------------------
 
 In order to dump additional data in WarpX (mostly for debugging purpose), run the simulation
-with parameters
+with parameters (``warpx.plot_rho`` only when using back-transformed diagnostics)
 
 ::
 

Docs/source/visualization/inputs.2d

@@ -9,7 +9,6 @@ amr.blocking_factor = 32
 amr.max_level = 0
 amr.plot_file = "plotfiles/plt"
 amr.plot_int = 0
-warpx.plot_rho = 1
 insitu.int = 2
 insitu.start = 0
 insitu.config = ez2d.xml

Docs/source/visualization/visualization.rst

@@ -56,7 +56,7 @@ quantity staggering
 E        n     n
 B        n     n
 j        n-1/2 n-1/2
-rho      n-1   n
+rho      n     n
 position n     n
 momentum n-1/2 n-1/2
 E part   n     n

Regression/Checksum/benchmarks_json/Langmuir_multi.json

@@ -19,7 +19,7 @@
     "jy": 6.087467475688316e+16,
     "jz": 6.087467475688315e+16,
     "part_per_cell": 524288.0,
-    "rho": 771603348.7307227
+    "rho": 702984843.3445112
   },
   "positrons": {
     "particle_cpu": 131072.0,
@@ -29,4 +29,4 @@
     "particle_position_y": 2.621440000000001,
     "particle_position_z": 2.6214400000000007
   }
-}
+}

Regression/Checksum/benchmarks_json/Langmuir_multi_nodal.json

@@ -19,7 +19,7 @@
     "jy": 5.803381905407015e+16,
     "jz": 5.803381905407016e+16,
     "part_per_cell": 524288.0,
-    "rho": 786759127.7574512
+    "rho": 720713352.6087718
   },
   "positrons": {
     "particle_cpu": 131072.0,
@@ -29,4 +29,4 @@
     "particle_position_y": 2.621440000000001,
     "particle_position_z": 2.6214399999999998
   }
-}
+}

Regression/Checksum/benchmarks_json/Langmuir_multi_psatd_single_precision.json

@@ -29,4 +29,4 @@
     "particle_position_y": 2.621439966217963,
     "particle_position_z": 2.6214399657132788
   }
-}
+}

Regression/Checksum/benchmarks_json/Langmuir_multi_single_precision.json

@@ -19,7 +19,7 @@
     "jy": 6.087430753764646e+16,
     "jz": 6.087430519585869e+16,
     "part_per_cell": 524288.0,
-    "rho": 771605196.0458984
+    "rho": 702985454.7099609
   },
   "positrons": {
     "particle_cpu": 131072.0,
@@ -29,4 +29,4 @@
     "particle_position_y": 2.621439958043453,
     "particle_position_z": 2.6214399581214423
   }
-}
+}

Regression/Checksum/benchmarks_json/Python_LaserAccelerationMR.json

@@ -20,6 +20,6 @@
     "jx": 1.928774119411483e+18,
     "jy": 69804337605269.28,
     "jz": 742686862695719.5,
-    "rho": 27689135301.61659
+    "rho": 27910721385.36522
   }
-}
+}

Source/Diagnostics/ComputeDiagFunctors/CMakeLists.txt

@@ -3,6 +3,7 @@ target_sources(WarpX
     CellCenterFunctor.cpp
     DivBFunctor.cpp
     DivEFunctor.cpp
+    RhoFunctor.cpp
     PartPerCellFunctor.cpp
     PartPerGridFunctor.cpp
     BackTransformFunctor.cpp

Source/Diagnostics/ComputeDiagFunctors/Make.package

@@ -3,6 +3,7 @@ CEXE_sources += PartPerCellFunctor.cpp
 CEXE_sources += PartPerGridFunctor.cpp
 CEXE_sources += DivBFunctor.cpp
 CEXE_sources += DivEFunctor.cpp
+CEXE_sources += RhoFunctor.cpp
 CEXE_sources += BackTransformFunctor.cpp
 
 VPATH_LOCATIONS   += $(WARPX_HOME)/Source/Diagnostics/ComputeDiagFunctors

Source/Diagnostics/ComputeDiagFunctors/RhoFunctor.H

@@ -0,0 +1,45 @@
+#ifndef WARPX_RHOFUNCTOR_H_
+#define WARPX_RHOFUNCTOR_H_
+
+#include "ComputeDiagFunctor.H"
+
+/**
+ * \brief Functor to compute charge density rho into mf_out
+ */
+class
+RhoFunctor final : public ComputeDiagFunctor
+{
+
+public:
+
+    /**
+     * \brief Constructor
+     *
+     * \param[in] lev level of MultiFab
+     * \param[in] crse_ratio coarsening ratio for interpolation of field values
+     *                       from simulation MultiFabs to the output MultiFab mf_dst
+     * \param[in] convertRZmodes2cartesian if true, all RZ modes are averaged into one component
+     * \param[in] ncomp optional number of component of source MultiFab mf_src
+     *                  to be cell-centered in output MultiFab mf_dst
+     */
+    RhoFunctor ( const int lev, const amrex::IntVect crse_ratio,
+                 bool convertRZmodes2cartesian=true, const int ncomp=1 );
+
+    /**
+     * \brief Compute rho directly into mf_dst
+     *
+     * \param[out] mf_dst output MultiFab where the result is written
+     * \param[in]  dcomp  first component of mf_dst in which cell-centered data are stored
+     */
+    virtual void operator() ( amrex::MultiFab& mf_dst, const int dcomp ) const override;
+
+private:
+
+    // Level on which source MultiFab mf_src is defined in RZ geometry
+    int const m_lev;
+
+    // Whether to average all modes into one component in RZ geometry
+    bool m_convertRZmodes2cartesian;
+};
+
+#endif // WARPX_RHOFUNCTOR_H_

Source/Diagnostics/ComputeDiagFunctors/RhoFunctor.cpp

@@ -0,0 +1,43 @@
+#include "WarpX.H"
+#include "RhoFunctor.H"
+#include "Utils/CoarsenIO.H"
+
+RhoFunctor::RhoFunctor ( const int lev, const amrex::IntVect crse_ratio,
+                         bool convertRZmodes2cartesian, const int ncomp )
+    : ComputeDiagFunctor(ncomp, crse_ratio), m_lev(lev),
+      m_convertRZmodes2cartesian(convertRZmodes2cartesian)
+{}
+
+void
+RhoFunctor::operator() ( amrex::MultiFab& mf_dst, const int dcomp ) const
+{
+    auto& warpx = WarpX::GetInstance();
+    auto& mypc  = warpx.GetPartContainer();
+
+    // Deposit charge density
+    std::unique_ptr<amrex::MultiFab> rho = mypc.GetChargeDensity( m_lev );
+
+#ifdef WARPX_DIM_RZ
+    if (m_convertRZmodes2cartesian) {
+        // In cylindrical geometry, sum real part of all modes of rho in
+        // temporary MultiFab mf_dst_stag, and cell-center it to mf_dst
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+            nComp()==1,
+            "The RZ averaging over modes must write into one single component");
+        amrex::MultiFab mf_dst_stag( rho->boxArray(), warpx.DistributionMap(m_lev), 1, rho->nGrowVect() );
+        // Mode 0
+        amrex::MultiFab::Copy( mf_dst_stag, *rho, 0, 0, 1, rho->nGrowVect() );
+        for (int ic=1 ; ic < rho->nComp() ; ic += 2) {
+            // Real part of all modes > 0
+            amrex::MultiFab::Add( mf_dst_stag, *rho, ic, 0, 1, rho->nGrowVect() );
+        }
+        CoarsenIO::Coarsen( mf_dst, mf_dst_stag, dcomp, 0, nComp(), 0, m_crse_ratio );
+    } else {
+        CoarsenIO::Coarsen( mf_dst, *rho, dcomp, 0, nComp(), 0, m_crse_ratio );
+    }
+#else
+    // In Cartesian geometry, coarsen and interpolate from temporary MultiFab rho
+    // to output diagnostic MultiFab mf_dst
+    CoarsenIO::Coarsen( mf_dst, *rho, dcomp, 0, nComp(), 0, m_crse_ratio );
+#endif
+}

Source/Diagnostics/Diagnostics.cpp

@@ -38,9 +38,11 @@ Diagnostics::BaseReadParameters ()
     if (!varnames_specified){
         m_varnames = {"Ex", "Ey", "Ez", "Bx", "By", "Bz", "jx", "jy", "jz"};
     }
-    // set plot_rho to true of the users requests it, so that
-    // rho is computed at each iteration.
-    if (WarpXUtilStr::is_in(m_varnames, "rho")) warpx.setplot_rho(true);
+    // If user requests rho with back-transformed diagnostics, we set plot_rho=true
+    // and compute rho at each iteration
+    if (WarpXUtilStr::is_in(m_varnames, "rho") && WarpX::do_back_transformed_diagnostics) {
+        warpx.setplot_rho(true);
+    }
     // Sanity check if user requests to plot F
     if (WarpXUtilStr::is_in(m_varnames, "F")){
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE(

Source/Diagnostics/FullDiagnostics.cpp

@@ -6,6 +6,7 @@
 #include "ComputeDiagFunctors/PartPerGridFunctor.H"
 #include "ComputeDiagFunctors/DivBFunctor.H"
 #include "ComputeDiagFunctors/DivEFunctor.H"
+#include "ComputeDiagFunctors/RhoFunctor.H"
 #include "FlushFormats/FlushFormat.H"
 #include "FlushFormats/FlushFormatPlotfile.H"
 #include "FlushFormats/FlushFormatCheckpoint.H"
@@ -48,7 +49,6 @@ FullDiagnostics::ReadParameters ()
 {
     // Read list of full diagnostics fields requested by the user.
     bool checkpoint_compatibility = BaseReadParameters();
-    auto & warpx = WarpX::GetInstance();
     amrex::ParmParse pp(m_diag_name);
     AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
         m_format == "plotfile" || m_format == "openpmd" ||
@@ -146,6 +146,9 @@ FullDiagnostics::AddRZModesToDiags (int lev)
         }
     }
 
+    // If rho is requested, all components will be written out
+    bool rho_requested = WarpXUtilStr::is_in( m_varnames, "rho" );
+
     // First index of m_all_field_functors[lev] where RZ modes are stored
     int icomp = m_all_field_functors[0].size();
     const std::array<std::string, 3> coord {"r", "theta", "z"};
@@ -156,6 +159,9 @@ FullDiagnostics::AddRZModesToDiags (int lev)
     if (divE_requested) {
         n_new_fields += 1;
     }
+    if (rho_requested) {
+        n_new_fields += 1;
+    }
     m_all_field_functors[lev].resize( m_all_field_functors[0].size() + n_new_fields );
     // E
     for (int dim=0; dim<3; dim++){
@@ -194,6 +200,12 @@ FullDiagnostics::AddRZModesToDiags (int lev)
         icomp += 1;
         AddRZModesToOutputNames(std::string("divE"), ncomp_multimodefab);
     }
+    // rho
+    if (rho_requested) {
+        m_all_field_functors[lev][icomp] = std::make_unique<RhoFunctor>(lev, m_crse_ratio, false, ncomp_multimodefab);
+        icomp += 1;
+        AddRZModesToOutputNames(std::string("rho"), ncomp_multimodefab);
+    }
     // Sum the number of components in input vector m_all_field_functors
     // and check that it corresponds to the number of components in m_varnames
     // and m_mf_output
@@ -342,13 +354,18 @@ FullDiagnostics::InitializeFieldFunctors (int lev)
         } else if ( m_varnames[comp] == "jz" ){
             m_all_field_functors[lev][comp] = std::make_unique<CellCenterFunctor>(warpx.get_pointer_current_fp(lev, 2), lev, m_crse_ratio);
         } else if ( m_varnames[comp] == "rho" ){
-            // rho_new is stored in component 1 of rho_fp when using PSATD
+            if ( WarpX::do_back_transformed_diagnostics ) {
 #ifdef WARPX_USE_PSATD
-            amrex::MultiFab* rho_new = new amrex::MultiFab(*warpx.get_pointer_rho_fp(lev), amrex::make_alias, 1, 1);
-            m_all_field_functors[lev][comp] = std::make_unique<CellCenterFunctor>(rho_new, lev, m_crse_ratio);
+                // rho_new is stored in component 1 of rho_fp when using PSATD
+                amrex::MultiFab* rho_new = new amrex::MultiFab(*warpx.get_pointer_rho_fp(lev), amrex::make_alias, 1, 1);
+                m_all_field_functors[lev][comp] = std::make_unique<CellCenterFunctor>(rho_new, lev, m_crse_ratio);
 #else
-            m_all_field_functors[lev][comp] = std::make_unique<CellCenterFunctor>(warpx.get_pointer_rho_fp(lev), lev, m_crse_ratio);
+                m_all_field_functors[lev][comp] = std::make_unique<CellCenterFunctor>(warpx.get_pointer_rho_fp(lev), lev, m_crse_ratio);
 #endif
+            }
+            else {
+                m_all_field_functors[lev][comp] = std::make_unique<RhoFunctor>(lev, m_crse_ratio);
+            }
         } else if ( m_varnames[comp] == "F" ){
             m_all_field_functors[lev][comp] = std::make_unique<CellCenterFunctor>(warpx.get_pointer_F_fp(lev), lev, m_crse_ratio);
         } else if ( m_varnames[comp] == "part_per_cell" ){

Source/WarpX.H

@@ -761,7 +761,7 @@ private:
 
     std::string restart_chkfile;
 
-    bool plot_rho           = false;
+    bool plot_rho = false;
 
     amrex::VisMF::Header::Version plotfile_headerversion  = amrex::VisMF::Header::Version_v1;
     amrex::VisMF::Header::Version slice_plotfile_headerversion  = amrex::VisMF::Header::Version_v1;

Source/WarpX.cpp

@@ -859,7 +859,7 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm
     current_fp[lev][1].reset( new MultiFab(amrex::convert(ba,jy_nodal_flag),dm,ncomps,ngJ));
     current_fp[lev][2].reset( new MultiFab(amrex::convert(ba,jz_nodal_flag),dm,ncomps,ngJ));
 
-    if (do_dive_cleaning || plot_rho)
+    if (do_dive_cleaning || (plot_rho && do_back_transformed_diagnostics))
     {
         rho_fp[lev].reset(new MultiFab(amrex::convert(ba,rho_nodal_flag),dm,2*ncomps,ngRho));
     }
@@ -968,7 +968,7 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm
         current_cp[lev][1].reset( new MultiFab(amrex::convert(cba,jy_nodal_flag),dm,ncomps,ngJ));
         current_cp[lev][2].reset( new MultiFab(amrex::convert(cba,jz_nodal_flag),dm,ncomps,ngJ));
 
-        if (do_dive_cleaning || plot_rho){
+        if (do_dive_cleaning || (plot_rho && do_back_transformed_diagnostics)) {
             rho_cp[lev].reset(new MultiFab(amrex::convert(cba,rho_nodal_flag),dm,2*ncomps,ngRho));
         }
         if (do_dive_cleaning)