Make laser output LASY compliant (#1046)

* laser output LASY compliant

* cleaning

* set attributes

* fix datatype

* fix datatype

* fix CI tests

* fix CI test and doc

* fix field_data laserEnvelope

* fix hasFieldOutput

* fix name collision

* add comment

* fix index bug

* merge dev

docs/source/run/parameters.rst

@@ -761,7 +761,7 @@ Laser parameters
 The laser profile is defined by :math:`a(x,y,z) = a_0 * \mathrm{exp}[-(x^2/w0_x^2 + y^2/w0_y^2 + z^2/L0^2)]`.
 The model implemented is the one from [C. Benedetti et al. Plasma Phys. Control. Fusion 60.1: 014002 (2017)].
 Unlike for ``beams`` and ``plasmas``, all the laser pulses are currently stored on the same array,
-which you can find in the output openPMD file as `laser_real` (for the real part of the envelope) and `laser_imag` for its imaginary part.
+which you can find in the output openPMD file as a complex array named `laserEnvelope`.
 Parameters starting with ``lasers.`` apply to all laser pulses, parameters starting with ``<laser name>`` apply to a single laser pulse.
 
 * ``lasers.names`` (list of `string`) optional (default `no_laser`)
@@ -897,7 +897,7 @@ Field diagnostics
     If ``rho`` is explicitly mentioned as ``field_data``, it is deposited by the plasma
     to be available as a diagnostic. Similarly if ``rho_<plasma name>`` is explicitly mentioned,
     the charge density of that plasma species will be separately available as a diagnostic.
-    When a laser pulse is used, the real and imaginary parts of the laser complex envelope are written in ``laser_real`` and ``laser_imag``, respectively.
+    When a laser pulse is used, the laser complex envelope ``laserEnvelope`` is available.
     The plasma proper density (n/gamma) is then also accessible via ``chi``.
 
 * ``<diag name> or diagnostic.patch_lo`` (3 `float`) optional (default `-infinity -infinity -infinity`)

examples/get_started/inputs_lwfa

@@ -11,8 +11,6 @@ my_constants.Lramp = 6.e-3
 hipace.verbose = 1
 hipace.dt = 10.*kp_inv/clight
 diagnostic.output_period = 10
-hipace.numprocs_x = 1
-hipace.numprocs_y = 1
 
 amr.max_level = 0
 

examples/get_started/inputs_pwfa

@@ -7,9 +7,6 @@ hipace.max_time = 0.3/clight
 diagnostic.output_period = 1
 hipace.verbose = 1
 
-hipace.numprocs_x = 1
-hipace.numprocs_y = 1
-
 hipace.depos_order_xy = 2
 hipace.dt = adaptive
 hipace.nt_per_betatron = 30

examples/laser/analysis_laser_vacuum.py

@@ -40,9 +40,8 @@
 A0 = np.zeros(len(ts1.iterations))
 Z = np.zeros(len(ts1.iterations))
 for iteration in ts1.iterations:
-    F1, m1 = ts1.get_field(iteration=iteration, field='laser_real')
-    F2, m2 = ts1.get_field(iteration=iteration, field='laser_imag')
-    a_abs = np.abs(F1+1j*F2)
+    F, m1 = ts1.get_field(iteration=iteration, field='laserEnvelope')
+    a_abs = np.abs(F)
     W0[iteration] = 2.*np.sqrt(np.sum(a_abs**2*m1.x**2)/np.sum(a_abs**2))
     A0[iteration] = 2.*np.max(a_abs)
     Z[iteration] = ts1.current_t*scc.c

src/Hipace.cpp

@@ -971,10 +971,7 @@ Hipace::FillFieldDiagnostics (const int lev, int islice)
 {
     for (auto& fd : m_diags.getFieldData()) {
         if (fd.m_level == lev && fd.m_has_field) {
-            m_fields.Copy(lev, islice, fd.m_geom_io, fd.m_F,
-                fd.m_F.box(), m_3D_geom[lev],
-                fd.m_comps_output_idx, fd.m_nfields,
-                fd.m_do_laser, m_multi_laser, fd.m_slice_dir);
+            m_fields.Copy(lev, islice, fd, m_3D_geom[lev], m_multi_laser);
         }
     }
 }
@@ -997,13 +994,13 @@ Hipace::WriteDiagnostics (const int step)
 #ifdef HIPACE_USE_OPENPMD
     if (m_diags.hasAnyFieldOutput(step, m_max_step, m_physical_time, m_max_time)) {
         m_openpmd_writer.WriteDiagnostics(m_diags.getFieldData(), m_multi_beam,
-                        m_physical_time, step, getDiagBeamNames(),
+                        m_multi_laser, m_physical_time, step, getDiagBeamNames(),
                         m_3D_geom, OpenPMDWriterCallType::fields);
     }
 
     if (m_diags.hasBeamOutput(step, m_max_step, m_physical_time, m_max_time)) {
         m_openpmd_writer.WriteDiagnostics(m_diags.getFieldData(), m_multi_beam,
-                        m_physical_time, step, getDiagBeamNames(),
+                        m_multi_laser, m_physical_time, step, getDiagBeamNames(),
                         m_3D_geom, OpenPMDWriterCallType::beams);
     }
 #else

src/diagnostics/Diagnostic.H

@@ -15,6 +15,19 @@
 
 #include <vector>
 
+// 16-byte alignment enables coalesced memory access which is significantly faster
+struct alignas(2*sizeof(amrex::Real)) AlignedComplex {
+    amrex::Real real;
+    amrex::Real imag;
+
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE constexpr
+    AlignedComplex& operator+= (const AlignedComplex& rhs) {
+        real += rhs.real;
+        imag += rhs.imag;
+        return *this;
+    }
+};
+
 
 /** \brief This struct holds data for one field diagnostic on one MR level */
 struct FieldDiagnosticData
@@ -32,17 +45,21 @@ struct FieldDiagnosticData
     amrex::IntVect m_diag_coarsen; /**< xyz coarsening ratio (positive) */
     bool m_do_laser {false}; /**< Whether to output the laser */
     int m_nfields; /**< Number of physical fields to write */
-    int m_nfields_with_laser; /**< Number of fields to write including laser */
     amrex::Vector<std::string> m_comps_output; /**< Component names to Write to output file */
     /** Component indexes to Write to output file */
     amrex::Gpu::DeviceVector<int> m_comps_output_idx;
     /** Vector over levels, all fields */
     amrex::FArrayBox m_F;
+    using complex_type = AlignedComplex;
+    /** FAB for laser */
+    amrex::BaseFab<complex_type> m_F_laser;
     amrex::Geometry m_geom_io; /**< Diagnostics geometry */
     bool m_has_field; /**< if there is field output to write */
     /** Number of iterations between consecutive output dumps.
      * Default is 0, meaning no output */
     int m_output_period = 0;
+    /** Name of the laser in input and output files */
+    std::string m_laser_io_name = "laserEnvelope";
 };
 
 
@@ -76,7 +93,7 @@ public:
                          int output_step, int max_step,
                          amrex::Real output_time, amrex::Real max_time)
     {
-        return  fd.m_nfields_with_laser > 0 &&
+        return  (fd.m_nfields > 0 || fd.m_do_laser) &&
             utils::doDiagnostics(fd.m_output_period, output_step, max_step,
                             output_time, max_time);
     }

src/diagnostics/Diagnostic.cpp

@@ -7,6 +7,7 @@
  */
 #include "Diagnostic.H"
 #include "Hipace.H"
+#include "utils/HipaceProfilerWrapper.H"
 #include <AMReX_ParmParse.H>
 
 Diagnostic::Diagnostic (int nlev)
@@ -125,30 +126,30 @@ Diagnostic::Initialize (const int lev, bool do_laser) {
     for (auto& fd : m_field_data) {
         amrex::ParmParse pp(fd.m_diag_name);
 
-        fd.m_do_laser = do_laser;
-
         queryWithParserAlt(pp, "field_data", fd.m_comps_output, ppd);
 
         amrex::Vector<std::string> all_field_comps{};
-        all_field_comps.reserve(Comps[WhichSlice::This].size());
+        all_field_comps.reserve(Comps[WhichSlice::This].size() + (do_laser ? 1 : 0));
         for (const auto& [comp, idx] : Comps[WhichSlice::This]) {
             all_field_comps.push_back(comp);
         }
+        if (do_laser) {
+            all_field_comps.push_back(fd.m_laser_io_name);
+        }
+
         if(fd.m_comps_output.empty()) {
             fd.m_comps_output = all_field_comps;
-        }
-        else {
-            for(const std::string& comp_name : fd.m_comps_output) {
-                if(comp_name == "all" || comp_name == "All") {
+        } else {
+            for (const std::string& comp_name : fd.m_comps_output) {
+                if (comp_name == "all" || comp_name == "All") {
                     fd.m_comps_output = all_field_comps;
                     break;
                 }
-                if(comp_name == "none" || comp_name == "None") {
+                if (comp_name == "none" || comp_name == "None") {
                     fd.m_comps_output.clear();
-                    fd.m_do_laser = false;
                     break;
                 }
-                if(Comps[WhichSlice::This].count(comp_name) == 0) {
+                if (std::find(all_field_comps.begin(), all_field_comps.end(), comp_name) ==  all_field_comps.end()) {
                     std::stringstream error_str{};
                     error_str << "Unknown field diagnostics component: " << comp_name <<"\nmust be "
                         << "'all', 'none' or a subset of:";
@@ -160,6 +161,16 @@ Diagnostic::Initialize (const int lev, bool do_laser) {
             }
         }
 
+        // remove laser from fd.m_comps_output because it is output separately
+        for (auto it = fd.m_comps_output.begin(); it != fd.m_comps_output.end();) {
+            if (*it == fd.m_laser_io_name) {
+                it = fd.m_comps_output.erase(it);
+                fd.m_do_laser = true;
+            } else {
+                ++it;
+            }
+        }
+
         fd.m_nfields = fd.m_comps_output.size();
 
         amrex::Gpu::PinnedVector<int> local_comps_output_idx(fd.m_nfields);
@@ -168,18 +179,13 @@ Diagnostic::Initialize (const int lev, bool do_laser) {
         }
         fd.m_comps_output_idx.assign(local_comps_output_idx.begin(), local_comps_output_idx.end());
 
-        fd.m_nfields_with_laser = fd.m_nfields;
-
-        if (fd.m_do_laser) {
-            fd.m_nfields_with_laser += 2;
-            fd.m_comps_output.push_back("laser_real");
-            fd.m_comps_output.push_back("laser_imag");
-        }
-
         if (m_field_data.size() != 1) {
             for (auto& comp_name : fd.m_comps_output) {
                 comp_name += "_" + fd.m_diag_name;
             }
+            if (fd.m_do_laser) {
+                fd.m_laser_io_name += "_" + fd.m_diag_name;
+            }
         }
     }
 
@@ -217,6 +223,8 @@ Diagnostic::ResizeFDiagFAB (const amrex::Box a_domain, const int lev,
                             amrex::Geometry const& geom, int output_step, int max_step,
                             amrex::Real output_time, amrex::Real max_time)
 {
+    HIPACE_PROFILE("Diagnostic::ResizeFDiagFAB()");
+
     AMREX_ALWAYS_ASSERT(m_initialized);
 
     for (auto& fd : m_field_data) {
@@ -272,8 +280,13 @@ Diagnostic::ResizeFDiagFAB (const amrex::Box a_domain, const int lev,
                          && hasFieldOutput(fd, output_step, max_step, output_time, max_time);
 
         if(fd.m_has_field) {
-            fd.m_F.resize(domain, fd.m_nfields_with_laser, amrex::The_Pinned_Arena());
+            fd.m_F.resize(domain, fd.m_nfields, amrex::The_Pinned_Arena());
             fd.m_F.setVal<amrex::RunOn::Host>(0);
+
+            if (fd.m_do_laser) {
+                fd.m_F_laser.resize(domain, 1, amrex::The_Pinned_Arena());
+                fd.m_F_laser.setVal<amrex::RunOn::Host>({0,0});
+            }
         }
     }
 }

src/diagnostics/OpenPMDWriter.H

@@ -11,6 +11,7 @@
 #include "diagnostics/Diagnostic.H"
 #include "particles/beam/MultiBeam.H"
 #include "particles/beam/BeamParticleContainer.H"
+#include "laser/MultiLaser.H"
 
 #include <AMReX_REAL.H>
 #include <AMReX_IntVect.H>
@@ -66,14 +67,11 @@ private:
 
     /** \brief writing openPMD field data
      *
-     * \param[in] fab the FArrayBox to dump
-     * \param[in] geom Geometry of the simulation, to get the cell size etc.
-     * \param[in] slice_dir direction of slicing. 0=x, 1=y, -1=no slicing (3D array)
-     * \param[in] varnames list of variable names for the fields (ExmBy, EypBx, Ey, ...)
+     * \param[in] fd field diagnostic data
+     * \param[in] a_multi_laser multi laser to get the central wavelength
      * \param[in,out] iteration openPMD iteration to which the data is written
      */
-    void WriteFieldData (amrex::FArrayBox const& fab, amrex::Geometry const& geom,
-                         const int slice_dir, const amrex::Vector< std::string > varnames,
+    void WriteFieldData (const FieldDiagnosticData& fd, const MultiLaser& a_multi_laser,
                          openPMD::Iteration iteration);
 
     /** Named Beam SoA attributes per particle as defined in BeamIdx
@@ -129,7 +127,7 @@ public:
      */
     void WriteDiagnostics (
         const amrex::Vector<FieldDiagnosticData>& field_diag, MultiBeam& a_multi_beam,
-        const amrex::Real physical_time, const int output_step,
+        const MultiLaser& a_multi_laser, const amrex::Real physical_time, const int output_step,
         const amrex::Vector< std::string > beamnames,
         amrex::Vector<amrex::Geometry> const& geom3D,
         const OpenPMDWriterCallType call_type);