add laser insitu diagnostics

docs/source/run/parameters.rst

@@ -934,9 +934,15 @@ 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], [Ez*jz_beam]``.
-Thereby, "[]" stands for averaging over all cells in the current slice.
+Thereby, "[]" stands for integrating 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, "[]" stands for integrating over all cells in the current slice, ``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
@@ -984,6 +990,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 field 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

@@ -1386,9 +1386,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,17 @@
 #include "Hipace.H"
 #include "utils/HipaceProfilerWrapper.H"
 #include "utils/DeprecatedInput.H"
+#include "utils/InsituUtil.H"
+#include "utils/TemplateUtil.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 +77,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 +184,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
@@ -1042,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);
+
+    TypeMultiplier<amrex::ReduceOps, amrex::ReduceOpMax, amrex::ReduceOpSum[m_insitu_nrp-1+m_insitu_ncp]> reduce_op;
+    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");