Add laser insitu diagnostics (#1060)

* add laser insitu diagnostics

* merge dev

* remove include

* add suggestions to documentation

docs/source/run/parameters.rst

@@ -924,21 +924,26 @@ 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 quantities, "[...]" stands for averaging over all particles in the current slice;
+for grid quantities, "[...]" stands for integrating over all cells in the current slice.
 
 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], [x*uy], [y*ux], [ux/uz], [uy/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.
+Thereby, "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.
 
 For the field in-situ diagnostics, the following quantities are calculated per slice and stored:
 ``[Ex^2], [Ey^2], [Ez^2], [Bx^2], [By^2], [Bz^2], [ExmBy^2], [EypBx^2], [jz_beam], [Ez*jz_beam]``.
-Thereby, "[]" stands for averaging over all cells in the current slice.
 These quantities can be used to calculate the energy stored in the fields.
 
+For the laser in-situ diagnostics, the following quantities are calculated per slice and stored:
+``max(|a|^2), [|a|^2], [|a|^2*x], [|a|^2*x*x], [|a|^2*y], [|a|^2*y*y], axis(a)``.
+Thereby, ``max(|a|^2)`` is the highest value of ``|a|^2`` in the current slice
+and ``axis(a)`` gives the complex value of the laser envelope, in the center of every slice.
+
 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
@@ -986,6 +991,12 @@ Use ``hipace/tools/read_insitu_diagnostics.py`` to read the files using this for
 * ``fields.insitu_file_prefix`` (`string`) optional (default ``"diags/field_insitu"``)
     Path of the field in-situ output. Must not be the same as `hipace.file_prefix`.
 
+* ``lasers.insitu_period`` (`int`) optional (default ``0``)
+    Period of the laser in-situ diagnostics. `0` means no laser in-situ diagnostics.
+
+* ``lasers.insitu_file_prefix`` (`string`) optional (default ``"diags/laser_insitu"``)
+    Path of the laser in-situ output. Must not be the same as `hipace.file_prefix`.
+
 Additional physics
 ------------------
 

src/Hipace.cpp

@@ -418,6 +418,7 @@ Hipace::Evolve ()
         m_fields.InSituWriteToFile(step, m_physical_time, m_3D_geom[0], m_max_step, m_max_time);
         m_multi_beam.InSituWriteToFile(step, m_physical_time, m_3D_geom[0], m_max_step, m_max_time);
         m_multi_plasma.InSituWriteToFile(step, m_physical_time, m_3D_geom[0], m_max_step, m_max_time);
+        m_multi_laser.InSituWriteToFile(step, m_physical_time, m_3D_geom[0], m_max_step, m_max_time);
 
         if (!m_explicit) {
             // averaging predictor corrector loop diagnostics
@@ -564,7 +565,10 @@ Hipace::SolveOneSlice (int islice, int step)
     // get field insitu diagnostics after all fields are computed & SALAME
     m_fields.InSituComputeDiags(step, m_physical_time, islice, m_3D_geom[0], m_max_step, m_max_time);
 
-    // copy fields to diagnostic array
+    // get laser insitu diagnostics
+    m_multi_laser.InSituComputeDiags(step, m_physical_time, islice, m_3D_geom[0], m_max_step, m_max_time);
+
+    // copy fields (and laser) to diagnostic array
     for (int lev=0; lev<current_N_level; ++lev) {
         FillFieldDiagnostics(lev, islice);
     }

src/fields/Fields.cpp

@@ -1388,9 +1388,9 @@ Fields::InSituWriteToFile (int step, amrex::Real time, const amrex::Geometry& ge
     ofs.close();
     // assert no file errors
 #ifdef HIPACE_USE_OPENPMD
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ofs, "Error while writing insitu beam diagnostics");
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ofs, "Error while writing insitu field diagnostics");
 #else
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ofs, "Error while writing insitu beam diagnostics. "
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ofs, "Error while writing insitu field diagnostics. "
         "Maybe the specified subdirectory does not exist");
 #endif
 

src/laser/MultiLaser.H

@@ -192,6 +192,27 @@ public:
      */
     void InitLaserSlice (const amrex::Geometry& geom, const int islice, const int comp);
 
+    /** Compute in-situ laser diagnostics of current slice, store in member variable
+     * \param[in] step current time step
+     * \param[in] time physical time
+     * \param[in] islice current slice, on which diags are computed.
+     * \param[in] geom3D Geometry of the problem
+     * \param[in] max_step maximum time step of simulation
+     * \param[in] max_time maximum time of simulation
+     */
+    void InSituComputeDiags (int step, amrex::Real time, int islice, const amrex::Geometry& geom3D,
+                             int max_step, amrex::Real max_time);
+
+    /** Dump in-situ reduced diagnostics to file.
+     * \param[in] step current time step
+     * \param[in] time physical time
+     * \param[in] geom3D Geometry object for the whole domain
+     * \param[in] max_step maximum time step of simulation
+     * \param[in] max_time maximum time of simulation
+     */
+    void InSituWriteToFile (int step, amrex::Real time, const amrex::Geometry& geom3D,
+                            int max_step, amrex::Real max_time);
+
     /** Get the central wavelength */
     amrex::Real GetLambda0 () const { return m_lambda0; }
 
@@ -256,6 +277,23 @@ private:
     cufftResult m_result_fwd;
     cufftResult m_result_bkw;
 #endif
+
+    // Data for in-situ diagnostics:
+    /** Number of real laser properties for in-situ per-slice reduced diagnostics. */
+    static constexpr int m_insitu_nrp = 6;
+    /** Number of real complex properties for in-situ per-slice reduced diagnostics. */
+    static constexpr int m_insitu_ncp = 1;
+    /** How often the insitu laser diagnostics should be computed and written
+     * Default is 0, meaning no output */
+    int m_insitu_period {0};
+    /** All per-slice real laser properties */
+    amrex::Vector<amrex::Real> m_insitu_rdata;
+    /** Sum of all per-slice real laser properties */
+    amrex::Vector<amrex::Real> m_insitu_sum_rdata;
+    /** All per-slice complex laser properties */
+    amrex::Vector<amrex::GpuComplex<amrex::Real>> m_insitu_cdata;
+    /** Prefix/path for the output files */
+    std::string m_insitu_file_prefix = "diags/laser_insitu";
 };
 
 #endif // MULTILASER_H_

src/laser/MultiLaser.cpp

@@ -12,12 +12,16 @@
 #include "Hipace.H"
 #include "utils/HipaceProfilerWrapper.H"
 #include "utils/DeprecatedInput.H"
+#include "utils/InsituUtil.H"
 #ifdef AMREX_USE_CUDA
 #  include "fields/fft_poisson_solver/fft/CuFFTUtils.H"
 #elif defined(AMREX_USE_HIP)
 #  include "fields/fft_poisson_solver/fft/RocFFTUtils.H"
 #endif
 #include "particles/particles_utils/ShapeFactors.H"
+#ifdef HIPACE_USE_OPENPMD
+#   include <openPMD/auxiliary/Filesystem.hpp>
+#endif
 
 #include <AMReX_GpuComplex.H>
 
@@ -72,6 +76,9 @@ MultiLaser::ReadParameters ()
         queryWithParser(pp, "openPMD_laser_name", m_file_envelope_name);
         queryWithParser(pp, "iteration", m_file_num_iteration);
     }
+
+    queryWithParser(pp, "insitu_period", m_insitu_period);
+    queryWithParser(pp, "insitu_file_prefix", m_insitu_file_prefix);
 }
 
 
@@ -176,6 +183,18 @@ MultiLaser::InitData (const amrex::BoxArray& slice_ba,
             amrex::ParallelDescriptor::Communicator());
 #endif
     }
+
+    if (m_insitu_period > 0) {
+#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 field insitu file prefix compared to the full diagnostics");
+#endif
+        // Allocate memory for in-situ diagnostics
+        m_insitu_rdata.resize(m_laser_geom_3D.Domain().length(2)*m_insitu_nrp, 0.);
+        m_insitu_sum_rdata.resize(m_insitu_nrp, 0.);
+        m_insitu_cdata.resize(m_laser_geom_3D.Domain().length(2)*m_insitu_ncp, 0.);
+    }
 }
 
 void
@@ -1043,3 +1062,161 @@ MultiLaser::InitLaserSlice (const amrex::Geometry& geom, const int islice, const
         }
     }
 }
+
+void
+MultiLaser::InSituComputeDiags (int step, amrex::Real time, int islice, const amrex::Geometry& geom3D,
+                                int max_step, amrex::Real max_time)
+{
+    if (!utils::doDiagnostics(m_insitu_period, step, max_step, time, max_time)) return;
+    HIPACE_PROFILE("MultiLaser::InSituComputeDiags()");
+
+    using namespace amrex::literals;
+    using Complex = amrex::GpuComplex<amrex::Real>;
+
+    AMREX_ALWAYS_ASSERT(m_insitu_rdata.size()>0 && m_insitu_sum_rdata.size()>0 &&
+                        m_insitu_cdata.size()>0);
+
+    const int nslices = geom3D.Domain().length(2);
+    const amrex::Real poff_x = GetPosOffset(0, geom3D, geom3D.Domain());
+    const amrex::Real poff_y = GetPosOffset(1, geom3D, geom3D.Domain());
+    const amrex::Real dx = geom3D.CellSize(0);
+    const amrex::Real dy = geom3D.CellSize(1);
+    const amrex::Real dxdydz = dx * dy * geom3D.CellSize(2);
+
+    const int xmid_lo = geom3D.Domain().smallEnd(0) + (geom3D.Domain().length(0) - 1) / 2;
+    const int xmid_hi = geom3D.Domain().smallEnd(0) + (geom3D.Domain().length(0)) / 2;
+    const int ymid_lo = geom3D.Domain().smallEnd(1) + (geom3D.Domain().length(1) - 1) / 2;
+    const int ymid_hi = geom3D.Domain().smallEnd(1) + (geom3D.Domain().length(1)) / 2;
+    const amrex::Real mid_factor = (xmid_lo == xmid_hi ? 1._rt : 0.5_rt)
+                                 * (ymid_lo == ymid_hi ? 1._rt : 0.5_rt);
+
+    amrex::TypeMultiplier<amrex::ReduceOps, amrex::ReduceOpMax, amrex::ReduceOpSum[m_insitu_nrp-1+m_insitu_ncp]> reduce_op;
+    amrex::TypeMultiplier<amrex::ReduceData, amrex::Real[m_insitu_nrp], Complex[m_insitu_ncp]> reduce_data(reduce_op);
+    using ReduceTuple = typename decltype(reduce_data)::Type;
+
+    for ( amrex::MFIter mfi(m_slices, DfltMfi); mfi.isValid(); ++mfi ) {
+        Array3<amrex::Real const> const arr = m_slices.const_array(mfi);
+        reduce_op.eval(
+            mfi.tilebox(), reduce_data,
+            [=] AMREX_GPU_DEVICE (int i, int j, int) -> ReduceTuple
+            {
+                using namespace WhichLaserSlice;
+                const amrex::Real areal = arr(i,j, n00j00_r);
+                const amrex::Real aimag = arr(i,j, n00j00_i);
+                const amrex::Real aabssq = abssq(areal, aimag);
+
+                const amrex::Real x = i * dx + poff_x;
+                const amrex::Real y = j * dy + poff_y;
+
+                const bool is_on_axis = (i==xmid_lo || i==xmid_hi) && (j==ymid_lo || j==ymid_hi);
+                const Complex aaxis{is_on_axis ? areal : 0._rt, is_on_axis ? aimag : 0._rt};
+
+                return {            // Tuple contains:
+                    aabssq,         // 0    max(|a|^2)
+                    aabssq,         // 1    [|a|^2]
+                    aabssq*x,       // 2    [|a|^2*x]
+                    aabssq*x*x,     // 3    [|a|^2*x*x]
+                    aabssq*y,       // 4    [|a|^2*y]
+                    aabssq*y*y,     // 5    [|a|^2*y*y]
+                    aaxis           // 6    axis(a)
+                };
+            });
+    }
+
+    ReduceTuple a = reduce_data.value();
+
+    amrex::constexpr_for<0, m_insitu_nrp>(
+        [&] (auto idx) {
+            if (idx.value == 0) {
+                m_insitu_rdata[islice + idx.value * nslices] = amrex::get<idx.value>(a);
+                m_insitu_sum_rdata[idx.value] =
+                    std::max(m_insitu_sum_rdata[idx.value], amrex::get<idx.value>(a));
+            } else {
+                m_insitu_rdata[islice + idx.value * nslices] = amrex::get<idx.value>(a)*dxdydz;
+                m_insitu_sum_rdata[idx.value] += amrex::get<idx.value>(a)*dxdydz;
+            }
+        }
+    );
+
+    amrex::constexpr_for<0, m_insitu_ncp>(
+        [&] (auto idx) {
+            m_insitu_cdata[islice + idx.value * nslices] =
+                amrex::get<m_insitu_nrp+idx.value>(a) * mid_factor;
+        }
+    );
+}
+
+void
+MultiLaser::InSituWriteToFile (int step, amrex::Real time, const amrex::Geometry& geom3D,
+                               int max_step, amrex::Real max_time)
+{
+    if (!utils::doDiagnostics(m_insitu_period, step, max_step, time, max_time)) return;
+    HIPACE_PROFILE("MultiLaser::InSituWriteToFile()");
+
+#ifdef HIPACE_USE_OPENPMD
+    // create subdirectory
+    openPMD::auxiliary::create_directories(m_insitu_file_prefix);
+#endif
+
+    // zero pad the rank number;
+    std::string::size_type n_zeros = 4;
+    std::string rank_num = std::to_string(amrex::ParallelDescriptor::MyProc());
+    std::string pad_rank_num = std::string(n_zeros-std::min(rank_num.size(), n_zeros),'0')+rank_num;
+
+    // open file
+    std::ofstream ofs{m_insitu_file_prefix + "/reduced_laser." + pad_rank_num + ".txt",
+        std::ofstream::out | std::ofstream::app | std::ofstream::binary};
+
+    const int nslices_int = geom3D.Domain().length(2);
+    const std::size_t nslices = static_cast<std::size_t>(nslices_int);
+    const int is_normalized_units = Hipace::m_normalized_units;
+
+    // specify the structure of the data later available in python
+    // avoid pointers to temporary objects as second argument, stack variables are ok
+    const amrex::Vector<insitu_utils::DataNode> all_data{
+        {"time"     , &time},
+        {"step"     , &step},
+        {"n_slices" , &nslices_int},
+        {"z_lo"     , &geom3D.ProbLo()[2]},
+        {"z_hi"     , &geom3D.ProbHi()[2]},
+        {"is_normalized_units", &is_normalized_units},
+        {"max(|a|^2)"     , &m_insitu_rdata[0], nslices},
+        {"[|a|^2]"        , &m_insitu_rdata[1*nslices], nslices},
+        {"[|a|^2*x]"      , &m_insitu_rdata[2*nslices], nslices},
+        {"[|a|^2*x*x]"    , &m_insitu_rdata[3*nslices], nslices},
+        {"[|a|^2*y]"      , &m_insitu_rdata[4*nslices], nslices},
+        {"[|a|^2*y*y]"    , &m_insitu_rdata[5*nslices], nslices},
+        {"axis(a)"        , &m_insitu_cdata[0], nslices},
+        {"integrated", {
+            {"max(|a|^2)"     , &m_insitu_sum_rdata[0]},
+            {"[|a|^2]"        , &m_insitu_sum_rdata[1]},
+            {"[|a|^2*x]"      , &m_insitu_sum_rdata[2]},
+            {"[|a|^2*x*x]"    , &m_insitu_sum_rdata[3]},
+            {"[|a|^2*y]"      , &m_insitu_sum_rdata[4]},
+            {"[|a|^2*y*y]"    , &m_insitu_sum_rdata[5]}
+        }}
+    };
+
+    if (ofs.tellp() == 0) {
+        // write JSON header containing a NumPy structured datatype
+        insitu_utils::write_header(all_data, ofs);
+    }
+
+    // write binary data according to datatype in header
+    insitu_utils::write_data(all_data, ofs);
+
+    // close file
+    ofs.close();
+    // assert no file errors
+#ifdef HIPACE_USE_OPENPMD
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ofs, "Error while writing insitu laser diagnostics");
+#else
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(ofs, "Error while writing insitu laser diagnostics. "
+        "Maybe the specified subdirectory does not exist");
+#endif
+
+    // reset arrays for insitu data
+    for (auto& x : m_insitu_rdata) x = 0.;
+    for (auto& x : m_insitu_sum_rdata) x = 0.;
+    for (auto& x : m_insitu_cdata) x = 0.;
+}

src/utils/InsituUtil.H

@@ -33,6 +33,7 @@ struct DataNode {
         std::string type_letter = "";
         if (std::is_integral_v<T> && std::is_signed_v<T>) type_letter = "i"; // int
         else if (std::is_integral_v<T> && std::is_unsigned_v<T>) type_letter = "u"; // unsiged int
+        else if (std::is_same_v<T, amrex::GpuComplex<amrex::Real>>) type_letter = "c"; // complex
         else if (std::is_floating_point_v<T>) type_letter = "f"; // float, double
         else amrex::Abort("DataNode (): unknown type T");