Plasma in-situ diags (#968)

* add plasma insitu diags

* add preliminary doc

* fix doxygen

* improve doc, add assert, add functions to tools

* remove typo

* add OpenMP support

* remove omp support

* Apply suggestions from code review

Co-authored-by: Maxence Thévenet <maxence.thevenet@desy.de>

* cleaning

* Update tools/read_insitu_diagnostics.py

Typo

* make eV consistent

---------

Co-authored-by: Maxence Thévenet <maxence.thevenet@desy.de>

docs/source/run/parameters.rst

@@ -491,40 +491,6 @@ which are valid only for certain beam types, are introduced further below under
     ``n_subcycles`` times with a time step of `dt/n_subcycles`. This can be used to improve accuracy
     in highly non-linear focusing fields.
 
-* ``<beam name> or beams.insitu_period`` (`int`) optional (default ``-1``)
-    Period of in-situ diagnostics, for computing the main per-slice beam quantities for the main
-    beam parameters (width, energy spread, emittance, etc.).
-    For this the following quantities are calculated per slice and stored:
-    ``sum(w), [x], [x^2], [y], [y^2], [z], [z^2], [ux], [ux^2], [uy], [uy^2], [uz], [uz^2], [x*ux], [y*uy], [z*uz], [ga], [ga^2], np``
-    where "[]" stands for averaging over all particles in the current slice,
-    "w" stands for weight, "ux" is the momentum in the x direction, "ga" is the Lorentz factor.
-    Averages and totals over all slices are also provided for convenience under the
-    respective ``average`` and ``total`` subcategories.
-
-    Additionally, some metadata is also available:
-    ``time, step, n_slices, charge, mass, z_lo, z_hi, normalized_density_factor``.
-    ``time`` and ``step`` refers to the physical time of the simulation and step number of the
-    current timestep.
-    ``n_slices`` equal to the number of slices in the zeta direction.
-    ``charge`` and ``mass`` relate to a single beam particle and are for example equal to the
-    electron charge and mass.
-    ``z_lo`` and ``z_hi`` are the lower and upper bounds of the z-axis of the simulation domain
-    specified in the input file and can be used to generate a z/zeta-axis for plotting.
-    ``normalized_density_factor`` is equal to ``dx * dy * dz`` in normalized units and 1 in
-    SI units. It can be used to convert ``sum(w)``, which specifies the beam density in normalized
-    units and beam weight an SI units, to the beam weight in both unit systems.
-
-    The data is written to a file at ``<insitu_file_prefix>/reduced_<beam name>.<MPI rank number>.txt``.
-    The in-situ diagnostics file format consists of a header part in ASCII containing a JSON object.
-    When this is parsed into Python it can be converted to a NumPy structured datatype.
-    The rest of the file, following immediately after the closing }, is in binary format and
-    contains all of the in-situ diagnostic along with some meta data. This part can be read using the
-    structured datatype of the first section.
-    Use ``hipace/tools/read_insitu_diagnostics.py`` to read the files using this format.
-
-* ``<beam name> or beams.insitu_file_prefix`` (`string`) optional (default ``"diags/insitu"``)
-    Path of the in-situ output.
-
 * ``<beam name>.do_salame`` (`bool`) optional (default `0`)
     Whether to use the SALAME algorithm [S. Diederichs et al., Phys. Rev. Accel. Beams 23, 121301 (2020)] to automatically flatten the accelerating field in the first time step. If turned on, the per-slice
     beam weight in the first time-step is adjusted such that the Ez field will be uniform in the beam.
@@ -686,8 +652,11 @@ Parameters starting with ``lasers.`` apply to all laser pulses, parameters start
 Diagnostic parameters
 ---------------------
 
+There are different types of diagnostics in HiPACE++. The standard diagnostics are compliant with the openPMD standard. The
+in-situ diagnostics allow for fast analysis of large beams or the plasma particles.
+
 * ``diagnostic.output_period`` (`integer`) optional (default `0`)
-    Output period for all diagnostics. Field or beam specific diagnostics can overwrite this parameter.
+    Output period for standard beam and field diagnostics. Field or beam specific diagnostics can overwrite this parameter.
     No output is given for ``diagnostic.output_period = 0``.
 
 Beam diagnostics
@@ -749,3 +718,52 @@ Field diagnostics
 
 * ``<diag name> or diagnostic.patch_hi`` (3 `float`) optional (default `infinity infinity infinity`)
     Upper limit for the diagnostic grid.
+
+In-situ diagnostics
+^^^^^^^^^^^^^^^^^^^
+
+Besides the standard diagnostics, fast in-situ diagnostics are available. They are most useful when beams with large numbers of particles are used, as the important moments can be calculated in-situ (during the simulation) to largely reduce the simulation's analysis.
+In-situ diagnostics compute slice quantities (1 number per quantity per longitudinal cell).
+For particle beams, they can be used to calculate the main characterizing beam parameters (width, energy spread, emittance, etc.), from which most common beam parameters (e.g. slice and projected emittance, etc.) can be computed. Additionally, the plasma particle properties (e.g, the temperature) can be calculated.
+
+For particle beams, the following quantities are calculated per slice and stored:
+``sum(w), [x], [x^2], [y], [y^2], [z], [z^2], [ux], [ux^2], [uy], [uy^2], [uz], [uz^2], [x*ux], [y*uy], [z*uz], [ga], [ga^2], np``.
+For plasma particles, the following quantities are calculated per slice and stored:
+``sum(w), [x], [x^2], [y], [y^2], [ux], [ux^2], [uy], [uy^2], [uz], [uz^2], [ga], [ga^2], np``.
+Thereby, "[]" stands for averaging over all particles in the current slice,
+"w" stands for weight, "ux" is the normalized momentum in the x direction, "ga" is the Lorentz factor.
+Averages and totals over all slices are also provided for convenience under the
+respective ``average`` and ``total`` subcategories.
+
+Additionally, some metadata is also available:
+``time, step, n_slices, charge, mass, z_lo, z_hi, normalized_density_factor``.
+``time`` and ``step`` refers to the physical time of the simulation and step number of the
+current timestep.
+``n_slices`` is the number of slices in the zeta direction.
+``charge`` and ``mass`` relate to a single particle and are for example equal to the
+electron charge and mass.
+``z_lo`` and ``z_hi`` are the lower and upper bounds of the z-axis of the simulation domain
+specified in the input file and can be used to generate a z/zeta-axis for plotting (note that they corresponds to mesh nodes, while the data is cell-centered).
+``normalized_density_factor`` is equal to ``dx * dy * dz`` in normalized units and 1 in
+SI units. It can be used to convert ``sum(w)``, which specifies the particle density in normalized
+units and particle weight in SI units, to the particle weight in both unit systems.
+
+The data is written to a file at ``<insitu_file_prefix>/reduced_<beam/plasma name>.<MPI rank number>.txt``.
+The in-situ diagnostics file format consists of a header part in ASCII containing a JSON object.
+When this is parsed into Python it can be converted to a NumPy structured datatype.
+The rest of the file, following immediately after the closing ``}``, is in binary format and
+contains all of the in-situ diagnostics along with some metadata. This part can be read using the
+structured datatype of the first section.
+Use ``hipace/tools/read_insitu_diagnostics.py`` to read the files using this format. Functions to calculate the most useful properties are also provided in that file.
+
+* ``<beam name> or beams.insitu_period`` (`int`) optional (default ``0``)
+    Period of the beam in-situ diagnostics. `0` means no beam in-situ diagnostics.
+
+* ``<beam name> or beams.insitu_file_prefix`` (`string`) optional (default ``"diags/insitu"``)
+    Path of the beam in-situ output. Must not be the same as `hipace.file_prefix`.
+
+* ``<plasma name> or plasmas.insitu_period`` (`int`) optional (default ``0``)
+    Period of the plasma in-situ diagnostics. `0` means no plasma in-situ diagnostics.
+
+* ``<plasma name> or plasmas.insitu_file_prefix`` (`string`) optional (default ``"plasma_diags/insitu"``)
+    Path of the plasma in-situ output. Must not be the same as `hipace.file_prefix`.

src/Hipace.cpp

@@ -459,6 +459,7 @@ Hipace::Evolve ()
         }
 
         m_multi_beam.InSituWriteToFile(step, m_physical_time, m_3D_geom[0]);
+        m_multi_plasma.InSituWriteToFile(step, m_physical_time, m_3D_geom[0]);
 
         // printing and resetting predictor corrector loop diagnostics
         if (m_verbose>=2 && !m_explicit) amrex::AllPrint() << "Rank " << rank <<
@@ -485,6 +486,7 @@ Hipace::SolveOneSlice (int islice, const int islice_local, int step)
     amrex::ParallelContext::push(m_comm_xy);
 
     m_multi_beam.InSituComputeDiags(step, islice, islice_local);
+    m_multi_plasma.InSituComputeDiags(step, islice);
 
     // Get this laser slice from the 3D array
     m_multi_laser.Copy(islice, false);

src/particles/beam/BeamParticleContainer.H

@@ -227,14 +227,13 @@ private:
      *      0,   1,     2,   3,     4,   5,     6,    7,      8,    9,     10,   11,     12,
      * sum(w), [x], [x^2], [y], [y^2], [z], [z^2], [ux], [ux^2], [uy], [uy^2], [uz], [uz^2],
      *
-     *     13,     14,     15,   16,     17, 18
-     * [x*ux], [y*uy], [z*uz], [ga], [ga^2], np
+     *     13,     14,     15,   16,     17
+     * [x*ux], [y*uy], [z*uz], [ga], [ga^2]
      * where [] means average over all particles within slice.
      * Per-slice emittance: sqrt( abs( ([x^2]-[x]^2) * ([ux^2]-[ux]^2) - ([x*ux]-[x][ux])^2 ) ).
      * Projected emittance: Same as above AFTER averaging all these quantities over slices.
      * Energy spread: sqrt([ga^2]-[ga]^2), and same as above.
      * Momentum spread: sqrt([uz^2]-[uz]^2), and same as above.
-     * Np: number of particles in this slice
      */
     amrex::Vector<amrex::Real> m_insitu_rdata;
     /** Per-slice int beam properties:
@@ -249,8 +248,9 @@ private:
     amrex::Vector<int> m_insitu_sum_idata;
     /** Prefix/path for the output files */
     std::string m_insitu_file_prefix = "diags/insitu";
-    /** How often the insitu beam diagnostics should be computed and written */
-    int m_insitu_period {-1};
+    /** How often the insitu beam diagnostics should be computed and written
+     * Default is 0, meaning no output */
+    int m_insitu_period {0};
 };
 
 #endif

src/particles/beam/BeamParticleContainer.cpp

@@ -189,7 +189,7 @@ BeamParticleContainer::InitData (const amrex::Geometry& geom)
 #ifdef HIPACE_USE_OPENPMD
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_insitu_file_prefix !=
             Hipace::GetInstance().m_openpmd_writer.m_file_prefix,
-            "Must choose a different insitu file prefix compared to the full diagnostics");
+            "Must choose a different beam insitu file prefix compared to the full diagnostics");
 #endif
         // Allocate memory for in-situ diagnostics
         m_nslices = geom.Domain().length(2);
@@ -312,8 +312,7 @@ void BeamParticleContainer::TagByLevel (const int current_N_level,
 bool
 BeamParticleContainer::doInSitu (int step)
 {
-    if (m_insitu_period <= 0) return false;
-    return step % m_insitu_period == 0;
+    return (m_insitu_period > 0 && step % m_insitu_period == 0);
 }
 
 void
@@ -326,7 +325,7 @@ BeamParticleContainer::InSituComputeDiags (int islice, int islice_local)
     AMREX_ALWAYS_ASSERT(m_insitu_rdata.size()>0 && m_insitu_idata.size()>0 &&
                         m_insitu_sum_rdata.size()>0 && m_insitu_sum_idata.size()>0);
 
-    PhysConst const phys_const = get_phys_const();
+    const PhysConst phys_const = get_phys_const();
     const amrex::Real clight_inv = 1.0_rt/phys_const.c;
     const amrex::Real clightsq_inv = 1.0_rt/(phys_const.c*phys_const.c);
 

src/particles/plasma/MultiPlasma.H

@@ -143,6 +143,13 @@ public:
      */
     void TagByLevel (const int current_N_level, amrex::Vector<amrex::Geometry> const& geom3D);
 
+    /** Compute reduced plasma diagnostics of current slice, store in member variable.
+     * \param[in] step time step of simulation
+     * \param[in] islice current slice, on which diags are computed.
+     */
+    void InSituComputeDiags (int step, int islice);
+    void InSituWriteToFile (int step, amrex::Real time, const amrex::Geometry& geom);
+
     int m_sort_bin_size {32}; /**< Tile size to sort plasma particles */
 
     /** Number of binary collisions */

src/particles/plasma/MultiPlasma.cpp

@@ -50,7 +50,7 @@ MultiPlasma::InitData (amrex::Vector<amrex::BoxArray> slice_ba,
     for (auto& plasma : m_all_plasmas) {
         // make it think there is only level 0
         plasma.SetParGDB(slice_gm[0], slice_dm[0], slice_ba[0]);
-        plasma.InitData();
+        plasma.InitData(gm[0]);
 
         if(plasma.m_can_ionize) {
             PlasmaParticleContainer* plasma_product = nullptr;
@@ -196,3 +196,23 @@ MultiPlasma::TagByLevel (const int current_N_level, amrex::Vector<amrex::Geometr
         plasma.TagByLevel(current_N_level, geom3D);
     }
 }
+
+void
+MultiPlasma::InSituComputeDiags (int step, int islice)
+{
+    for (auto& plasma : m_all_plasmas) {
+        if (plasma.doInSitu(step)) {
+            plasma.InSituComputeDiags(islice);
+        }
+    }
+}
+
+void
+MultiPlasma::InSituWriteToFile (int step, amrex::Real time, const amrex::Geometry& geom)
+{
+    for (auto& plasma : m_all_plasmas) {
+        if (plasma.doInSitu(step)) {
+            plasma.InSituWriteToFile(step, time, geom);
+        }
+    }
+}

src/particles/plasma/PlasmaParticleContainer.H

@@ -64,8 +64,9 @@ public:
     void ReadParameters ();
 
     /** Allocate data for the beam particles and initialize particles with requested beam profile
+    * \param[in] geom Geometry object for the whole domain
      */
-    void InitData ();
+    void InitData (const amrex::Geometry& geom);
 
     /** Initialize 1 xy slice of particles, with fixed number of particles per cell
      *
@@ -137,6 +138,18 @@ public:
     /** Returns name of the plasma */
     const std::string& GetName () const {return m_name;}
 
+    /** Compute in-situ plasma diagnostics of current slice, store in member variable
+     * \param[in] islice current slice, on which diags are computed.
+     */
+    void InSituComputeDiags (int islice);
+    /** Dump in-situ reduced diagnostics to file.
+     * \param[in] step current time step
+     * \param[in] time physical time
+     * \param[in] geom Geometry object for the whole domain
+     */
+    void InSituWriteToFile (int step, amrex::Real time, const amrex::Geometry& geom);
+    bool doInSitu (int step);
+
     amrex::Parser m_parser; /**< owns data for m_density_func */
     amrex::ParserExecutor<3> m_density_func; /**< Density function for the plasma */
     amrex::Real m_min_density {0.}; /**< minimal density at which particles are injected */
@@ -174,6 +187,35 @@ public:
 
 private:
     std::string m_name; /**< name of the species */
+    int m_nslices; /**< number of z slices of the domain */
+    /** Number of real plasma properties for in-situ per-slice reduced diagnostics. */
+    int m_insitu_nrp {13};
+    /** Number of int plasma properties for in-situ per-slice reduced diagnostics. */
+    int m_insitu_nip {1};
+    /** Per-slice real plasma properties:
+     *      0,   1,     2,   3,     4,    5,      6,    7,      8,      9,     10,   11,     12
+     * sum(w), [x], [x^2], [y], [y^2], [ux], [ux^2], [uy], [uy^2], [x*ux], [y*uy], [ga], [ga^2]
+     * where [] means average over all particles within slice.
+     * Per-slice emittance: sqrt( abs( ([x^2]-[x]^2) * ([ux^2]-[ux]^2) - ([x*ux]-[x][ux])^2 ) ).
+     * Projected emittance: Same as above AFTER averaging all these quantities over slices.
+     * Energy spread: sqrt([ga^2]-[ga]^2), and same as above.
+     */
+    amrex::Vector<amrex::Real> m_insitu_rdata;
+    /** Per-slice int plasma properties:
+     *  0
+     * Np
+     * Np: number of particles in this slice
+     */
+    amrex::Vector<int> m_insitu_idata;
+    /** Sum of all per-slice real plasma properties */
+    amrex::Vector<amrex::Real> m_insitu_sum_rdata;
+    /** Sum of all per-slice int plasma properties */
+    amrex::Vector<int> m_insitu_sum_idata;
+    /** Prefix/path for the output files */
+    std::string m_insitu_file_prefix = "diags/plasma_insitu";
+    /** How often the insitu plasma diagnostics should be computed and written
+     * Default is 0, meaning no output */
+    int m_insitu_period {0};
 };
 
 /** \brief Iterator over boxes in a particle container */