Encapsulate field IO, first step to multiple slice IO

doc/source/run/parameters.rst

@@ -39,14 +39,6 @@ General parameters
     | Output period. No output is given for `hipace.output_period = -1`.
     | **Warning:** `hipace.output_period = 0` will make the simulation crash.
 
-* ``hipace.output_slice`` (`bool`) optional (default `0`)
-    | Gives only a 2D slice output in the XZ-plane. The output is averaged over
-      the two central grid points of the y-axis.
-
-* ``hipace.3d_on_host`` (`bool`) optional (default `0`)
-    | Allocates the 3D arrays of the fields for I/O on the host, and not the device.
-      This enables to run simulations, where the 3D array exceeds the GPU memory.
-
 * ``hipace.beam_injection_cr`` (`integer`) optional (default `1`)
     | Using a temporary coarsed grid for beam particle injection for a fixed particle-per-cell beam.
       For very high-resolution simulations, where the number of grid points (`nx*ny*nz`)

examples/beam_in_vacuum/inputs_SI

@@ -36,3 +36,5 @@ beam2.density = 1.4119793504295784e23 # at this density, 1/kp = 10um
 beam2.u_mean = 0. 0. 1.e3
 beam2.u_std = 0. 0. 0.
 beam2.ppc = 2 2 1
+
+diagnostic.diag_type = xyz

examples/beam_in_vacuum/inputs_normalized

@@ -27,3 +27,5 @@ beam.density = 1.0
 beam.u_mean = 0. 0. 1.e3
 beam.u_std = 0. 0. 0.
 beam.ppc = 2 2 1
+
+diagnostic.diag_type = xyz

examples/blowout_wake/analysis_slice_IO.py

@@ -22,33 +22,57 @@
                                      left_edge=ds1.domain_left_edge,
                                      dims=ds1.domain_dimensions)
 F_full = all_data_level_0[field].v.squeeze()
-F_full_sliced = (F_full[:,F_full.shape[1]//2,:].squeeze() +
-                  F_full[:,F_full.shape[1]//2-1,:].squeeze())/2.
-ds2 = AMReXDataset('slice_io')
+F_full_xz = (F_full[:,F_full.shape[1]//2,:].squeeze() +
+             F_full[:,F_full.shape[1]//2-1,:].squeeze())/2.
+F_full_yz = (F_full[F_full.shape[0]//2,:,:].squeeze() +
+             F_full[F_full.shape[0]//2-1,:,:].squeeze())/2.
+
+ds2 = AMReXDataset('slice_io_xz')
 ad = ds2.all_data()
-F_slice = ad[field].reshape(ds2.domain_dimensions).v.squeeze()
+F_slice_xz = ad[field].reshape(ds2.domain_dimensions).v.squeeze()
+
+ds3 = AMReXDataset('slice_io_yz')
+ad = ds3.all_data()
+F_slice_yz = ad[field].reshape(ds3.domain_dimensions).v.squeeze()
 
 if do_plot:
-    plt.figure(figsize=(12,4))
-    plt.subplot(131)
-    plt.title('full')
-    plt.imshow(F_full_sliced)
+    plt.figure(figsize=(12,8))
+    plt.subplot(231)
+    plt.title('full_xz')
+    plt.imshow(F_full_xz)
+    plt.colorbar()
+    plt.subplot(232)
+    plt.title('slice xz')
+    plt.imshow(F_slice_xz)
+    plt.colorbar()
+    plt.subplot(233)
+    plt.title('difference')
+    plt.imshow(F_slice_xz-F_full_xz)
+    plt.colorbar()
+
+    plt.subplot(234)
+    plt.title('full_yz')
+    plt.imshow(F_full_yz)
     plt.colorbar()
-    plt.subplot(132)
-    plt.title('slice')
-    plt.imshow(F_slice)
+    plt.subplot(235)
+    plt.title('slice yz')
+    plt.imshow(F_slice_yz)
     plt.colorbar()
-    plt.subplot(133)
+    plt.subplot(236)
     plt.title('difference')
-    plt.imshow(F_slice-F_full_sliced)
+    plt.imshow(F_slice_yz-F_full_yz)
     plt.colorbar()
     plt.tight_layout()
     plt.savefig("image.pdf", bbox_inches='tight')
 
-error = np.max(np.abs(F_slice-F_full_sliced)) / np.max(np.abs(F_full_sliced))
+error_xz = np.max(np.abs(F_slice_xz-F_full_xz)) / np.max(np.abs(F_full_xz))
+error_yz = np.max(np.abs(F_slice_yz-F_full_yz)) / np.max(np.abs(F_full_yz))
 
 print("F_full.shape", F_full.shape)
-print("F_slice.shape", F_slice.shape)
-print("error", error)
+print("F_slice_xz.shape", F_slice_xz.shape)
+print("F_slice_yz.shape", F_slice_yz.shape)
+print("error_xz", error_xz)
+print("error_yz", error_yz)
 
-assert(error < 1.e-14)
+assert(error_xz < 1.e-14)
+assert(error_yz < 1.e-14)

examples/blowout_wake/inputs_SI

@@ -36,3 +36,5 @@ beam.ppc = 1 1 1
 plasma.density =  2.8239587008591567e23 # at this density, 1/kp = 10um
 plasma.ppc = 1 1
 plasma.u_mean = 0.0 0.0 0.
+
+diagnostic.diag_type = xyz

examples/blowout_wake/inputs_normalized

@@ -37,3 +37,5 @@ beam.ppc = 1 1 1
 plasma.density = 1.
 plasma.ppc = 1 1
 plasma.u_mean = 0.0 0.0 0.
+
+diagnostic.diag_type = xyz

examples/gaussian_weight/inputs_SI

@@ -26,3 +26,5 @@ beam.radius = 1.
 beam.u_mean = 0. 0. 1.e3
 beam.u_std = 0. 0. 0.
 beam.ppc = 2 2 1
+
+diagnostic.diag_type = xyz

examples/gaussian_weight/inputs_normalized

@@ -26,3 +26,5 @@ beam.zmax =  20.
 beam.radius = 40.
 beam.u_mean = 0. 0. 1.e3
 beam.u_std = 3. 4. 0.
+
+diagnostic.diag_type = xyz

examples/linear_wake/inputs_SI

@@ -1,6 +1,5 @@
 amr.n_cell = 32 32 200
 
-hipace.3d_on_host = 1
 hipace.output_plasma = 1
 
 amr.blocking_factor = 2
@@ -33,3 +32,5 @@ beam.ppc = 1 1 1
 plasma.density =  2.8239587008591567e23 # at this density, 1/kp = 10um
 plasma.ppc = 1 1
 plasma.u_mean = 0.0 0.0 0.
+
+diagnostic.diag_type = xyz

examples/linear_wake/inputs_normalized

@@ -1,7 +1,6 @@
 amr.n_cell = 32 32 200
 
 hipace.normalized_units=1
-hipace.3d_on_host = 1
 
 hipace.output_period = 1
 
@@ -34,3 +33,5 @@ beam.ppc = 1 1 1
 plasma.density = 1.
 plasma.ppc = 1 1
 plasma.u_mean = 0.0 0.0 0.
+
+diagnostic.diag_type = xyz

src/Hipace.H

@@ -193,13 +193,8 @@ public:
     /** Mixing factor between the transverse B field iterations in the predictor corrector loop
      */
     static amrex::Real m_predcorr_B_mixing_factor;
-    /** whether the 3D array stays in host memory (Pinned memory) or is allowed
-     * to be in device memory (allocated in Managed memory). */
-    static bool m_3d_on_host;
     /** Whether to call amrex::Gpu::synchronize() around all profiler region */
     static bool m_do_device_synchronize;
-    /** whether to output the central y=0 slice (if true), or the whole 3D array (if false) */
-    static bool m_slice_F_xz;
     /** Whether to dump plasma particles */
     static bool m_output_plasma;
     /** How much the box is coarsened for beam injection, to avoid exceeding max int in cell count.

src/Hipace.cpp

@@ -31,9 +31,7 @@ int Hipace::m_depos_order_z = 0;
 amrex::Real Hipace::m_predcorr_B_error_tolerance = 4e-2;
 int Hipace::m_predcorr_max_iterations = 30;
 amrex::Real Hipace::m_predcorr_B_mixing_factor = 0.05;
-bool Hipace::m_3d_on_host = false;
 bool Hipace::m_do_device_synchronize = false;
-bool Hipace::m_slice_F_xz = false;
 bool Hipace::m_output_plasma = false;
 int Hipace::m_beam_injection_cr = 1;
 amrex::Real Hipace::m_external_focusing_field_strength = 0.;
@@ -76,9 +74,7 @@ Hipace::Hipace () :
     pph.query("predcorr_max_iterations", m_predcorr_max_iterations);
     pph.query("predcorr_B_mixing_factor", m_predcorr_B_mixing_factor);
     pph.query("output_period", m_output_period);
-    pph.query("output_slice", m_slice_F_xz);
     pph.query("beam_injection_cr", m_beam_injection_cr);
-    pph.query("3d_on_host", m_3d_on_host);
     m_numprocs_z = amrex::ParallelDescriptor::NProcs() / (m_numprocs_x*m_numprocs_y);
     AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_numprocs_x*m_numprocs_y*m_numprocs_z
                                      == amrex::ParallelDescriptor::NProcs(),
@@ -307,7 +303,7 @@ Hipace::Evolve ()
         const amrex::Vector<int> index_array = fields.IndexArray();
         for (auto it = index_array.rbegin(); it != index_array.rend(); ++it)
         {
-            const amrex::Box& bx = fields.box(*it);
+            const amrex::Box& bx = boxArray(lev)[*it];
             amrex::Vector<amrex::DenseBins<BeamParticleContainer::ParticleType>> bins;
             bins = m_multi_beam.findParticlesInEachSlice(lev, *it, bx, geom[lev]);
 
@@ -391,7 +387,7 @@ Hipace::SolveOneSlice (int islice, int lev, amrex::Vector<amrex::DenseBins<BeamP
     // Push beam particles
     m_multi_beam.AdvanceBeamParticlesSlice(m_fields, geom[lev], lev, islice, bins);
 
-    m_fields.Copy(lev, islice, FieldCopyType::StoF, 0, 0, FieldComps::nfields);
+    m_fields.FillDiagnostics(lev, islice);
 
     m_fields.ShiftSlices(lev);
 
@@ -747,6 +743,7 @@ void
 Hipace::WriteDiagnostics (int output_step, bool force_output)
 {
     HIPACE_PROFILE("Hipace::WriteDiagnostics()");
+    constexpr int lev = 0;
 
     // Dump before first and after last step, and every m_output_period steps in-between
     if (m_output_period < 0 ||
@@ -763,30 +760,17 @@ Hipace::WriteDiagnostics (int output_step, bool force_output)
         {"ExmBy", "EypBx", "Ez", "Bx", "By", "Bz", "jx", "jy", "jz", "rho", "Psi"};
 
     amrex::Vector<std::string> rfs;
-    amrex::Vector<amrex::Geometry> geom_io = Geom();
-
-    constexpr int lev = 0; // we only have one level
-    constexpr int idim = 1;
-    amrex::RealBox prob_domain = Geom(lev).ProbDomain();
-    amrex::Box domain = Geom(lev).Domain();
-    // Define slice box
-    int const icenter = domain.length(idim)/2;
-    domain.setSmall(idim, icenter);
-    domain.setBig(idim, icenter);
-    if (m_slice_F_xz){
-        geom_io[lev] = amrex::Geometry(domain, &prob_domain, Geom(lev).Coord());
-    }
 
 #ifdef HIPACE_USE_OPENPMD
-    m_openpmd_writer.WriteDiagnostics(m_fields, m_multi_beam, geom_io[lev], m_physical_time,
-                                      output_step,  lev, m_slice_F_xz, varnames);
+    m_openpmd_writer.WriteDiagnostics(m_fields.getDiagF(), m_multi_beam, m_fields.getDiagGeom(),
+                                      m_physical_time, output_step,  lev, m_fields.getDiagSliceDir(), varnames);
 #else
     constexpr int nlev = 1;
     const amrex::IntVect local_ref_ratio {1, 1, 1};
 
     amrex::WriteMultiLevelPlotfile(
-        filename, nlev, amrex::GetVecOfConstPtrs(m_fields.getF()), varnames,
-        geom_io, m_physical_time, {output_step}, {local_ref_ratio},
+        filename, nlev, amrex::GetVecOfConstPtrs(m_fields.getDiagF()), varnames,
+        m_fields.getDiagGeom(), m_physical_time, {output_step}, {local_ref_ratio},
         "HyperCLaw-V1.1", "Level_", "Cell", rfs);
 
     // Write beam particles

src/diagnostics/CMakeLists.txt

@@ -1,4 +1,5 @@
 target_sources(HiPACE
   PRIVATE
     OpenPMDWriter.cpp
+    FieldDiagnostic.cpp
 )

src/diagnostics/FieldDiagnostic.H

@@ -0,0 +1,53 @@
+#ifndef FIELDDIAGNOSTIC_H_
+#define FIELDDIAGNOSTIC_H_
+
+#include <AMReX_MultiFab.H>
+
+/** type of diagnostics: full xyz array or xz slice or yz slice */
+enum struct DiagType{xyz, xz, yz};
+
+/** \brief This class holds data for 1 diagnostics (full or slice) */
+class FieldDiagnostic
+{
+
+public:
+
+    /** \brief Constructor */
+    explicit FieldDiagnostic (int nlev);
+
+    /** \brief allocate arrays of this MF
+     *
+     * \param[in] lev MR level
+     * \param[in] ba BoxArray of the full simulation domain
+     * \param[in] nfields number of field components
+     * \param[in] dm DistributionMapping of the full simulation domain
+     * \param[in] geom geometry of the full simulation domain
+     */
+    void AllocData (int lev, const amrex::BoxArray& ba, int nfields, const amrex::DistributionMapping& dm, amrex::Geometry const& geom);
+
+    /** \brief return the main diagnostics multifab */
+    amrex::Vector<amrex::MultiFab>& getF () { return m_F; };
+
+    /** \brief return the main diagnostics multifab
+     *
+     * \param[in] lev MR level
+     */
+    amrex::MultiFab& getF (int lev) {return m_F[lev];};
+
+    /** return the diagnostics geometry */
+    amrex::Vector<amrex::Geometry>& getGeom () { return m_geom_io; };
+
+    /** return slice direction of the diagnostics */
+    int sliceDir () {return m_slice_dir;};
+
+private:
+
+    /** Vector over levels, all fields */
+    amrex::Vector<amrex::MultiFab> m_F;
+    DiagType m_diag_type; /**< Type of diagnostics (xyz xz yz) */
+    int m_slice_dir; /**< Slicing direction */
+    int m_nfields; /**< Number of physical fields to write */
+    amrex::Vector<amrex::Geometry> m_geom_io; /**< Diagnostics geometry */
+};
+
+#endif // FIELDDIAGNOSTIC_H_

src/diagnostics/FieldDiagnostic.cpp

@@ -0,0 +1,61 @@
+#include "FieldDiagnostic.H"
+
+#include <AMReX_ParmParse.H>
+
+FieldDiagnostic::FieldDiagnostic (int nlev)
+    : m_F(nlev),
+      m_geom_io(nlev)
+{
+    amrex::ParmParse ppd("diagnostic");
+    std::string str_type;
+    ppd.get("diag_type", str_type);
+    if        (str_type == "xyz"){
+        m_diag_type = DiagType::xyz;
+        m_slice_dir = -1;
+    } else if (str_type == "xz") {
+        m_diag_type = DiagType::xz;
+        m_slice_dir = 1;
+    } else if (str_type == "yz") {
+        m_diag_type = DiagType::yz;
+        m_slice_dir = 0;
+    } else {
+        amrex::Abort("Unknown diagnostics type: must be xyz, xz or yz.");
+    }
+}
+
+void
+FieldDiagnostic::AllocData (int lev, const amrex::BoxArray& ba, int nfields, const amrex::DistributionMapping& dm, amrex::Geometry const& geom)
+{
+    m_nfields = nfields;
+    // Create a xz slice BoxArray
+    amrex::BoxList F_boxes;
+    if (m_slice_dir >= 0){
+        for (int i = 0; i < ba.size(); ++i){
+            amrex::Box bx = ba[i];
+            // Flatten the box down to 1 cell in the approprate direction.
+            bx.setSmall(m_slice_dir, ba[i].length(m_slice_dir)/2);
+            bx.setBig  (m_slice_dir, ba[i].length(m_slice_dir)/2);
+            // Note: the MR is still cell-centered, although the data will be averaged to nodal.
+            F_boxes.push_back(bx);
+        }
+    }
+    amrex::BoxArray F_slice_ba(std::move(F_boxes));
+    // m_F is defined on F_ba, the full or the slice BoxArray
+    amrex::BoxArray F_ba = m_slice_dir >= 0 ? F_slice_ba : ba;
+    // Only xy slices need guard cells, there is no deposition to/gather from the output array F.
+    amrex::IntVect nguards_F = amrex::IntVect(0,0,0);
+    // The Arena uses pinned memory.
+    m_F[lev].define(F_ba, dm, m_nfields, nguards_F,
+                    amrex::MFInfo().SetArena(amrex::The_Pinned_Arena()));
+
+    m_geom_io[lev] = geom;
+    amrex::RealBox prob_domain = geom.ProbDomain();
+    amrex::Box domain = geom.Domain();
+    // Define slice box
+    if (m_slice_dir >= 0){
+        int const icenter = domain.length(m_slice_dir)/2;
+        domain.setSmall(m_slice_dir, icenter);
+        domain.setBig(m_slice_dir, icenter);
+        m_geom_io[lev] = amrex::Geometry(domain, &prob_domain, geom.Coord());
+    }
+}