Allow for arbitrary sized diagnostics grid (#871)

* Allow for arbitrary sized diagnostics grid

* merge dev

* treat lo and hi seperately

* add CI test 😒

* remove some dev comments

docs/source/run/parameters.rst

@@ -630,3 +630,9 @@ Diagnostic parameters
     ``none`` or a subset of ``beams.names``.
     **Note:** The option ``none`` only suppressed the output of the beam data. To suppress any
     output, please use ``hipace.output_period = -1``.
+
+* ``diagnostic.patch_lo`` (3 `float`) optional (default `-infinity -infinity -infinity`)
+    Lower limit for the diagnostic grid.
+
+* ``diagnostic.patch_hi`` (3 `float`) optional (default `infinity infinity infinity`)
+    Upper limit for the diagnostic grid.

examples/blowout_wake/analysis_slice_IO.py

@@ -30,51 +30,96 @@
              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.
+F_full_cut_xy = F_full[20:45,:,50].squeeze()
+F_full_cut_xz = (F_full[32:49,43,:].squeeze() + F_full[32:49,44,:].squeeze())/2.
 
 ts2 = OpenPMDTimeSeries('slice_io_xz')
 F_slice_xz = ts2.get_field(field=field, iteration=ts2.iterations[-1])[0].transpose()
 
 ts3 = OpenPMDTimeSeries('slice_io_yz')
 F_slice_yz = ts3.get_field(field=field, iteration=ts3.iterations[-1])[0].transpose()
 
+ts4 = OpenPMDTimeSeries('slice_io_cut_xy')
+F_slice_cut_xy = ts4.get_field(field=field, iteration=ts4.iterations[-1])[0].squeeze().transpose()
+
+ts5 = OpenPMDTimeSeries('slice_io_cut_xz')
+F_slice_cut_xz = ts5.get_field(field=field, iteration=ts5.iterations[-1])[0].transpose()
+
 if do_plot:
-    plt.figure(figsize=(12,8))
-    plt.subplot(231)
+    plt.figure(figsize=(12,16))
+    plt.subplot(431)
     plt.title('full_xz')
     plt.imshow(F_full_xz)
     plt.colorbar()
-    plt.subplot(232)
+    plt.subplot(432)
     plt.title('slice xz')
     plt.imshow(F_slice_xz)
     plt.colorbar()
-    plt.subplot(233)
+    plt.subplot(433)
     plt.title('difference')
     plt.imshow(F_slice_xz-F_full_xz)
     plt.colorbar()
 
-    plt.subplot(234)
+    plt.subplot(434)
     plt.title('full_yz')
     plt.imshow(F_full_yz)
     plt.colorbar()
-    plt.subplot(235)
+    plt.subplot(435)
     plt.title('slice yz')
     plt.imshow(F_slice_yz)
     plt.colorbar()
-    plt.subplot(236)
+    plt.subplot(436)
     plt.title('difference')
     plt.imshow(F_slice_yz-F_full_yz)
     plt.colorbar()
     plt.tight_layout()
+
+    plt.subplot(437)
+    plt.title('full_xy')
+    plt.imshow(F_full_cut_xy)
+    plt.colorbar()
+    plt.subplot(438)
+    plt.title('cut slice xy')
+    plt.imshow(F_slice_cut_xy)
+    plt.colorbar()
+    plt.subplot(439)
+    plt.title('difference')
+    plt.imshow(F_slice_cut_xy-F_full_cut_xy)
+    plt.colorbar()
+    plt.tight_layout()
+
+    plt.subplot(4, 3, 10)
+    plt.title('full_xz')
+    plt.imshow(F_full_cut_xz)
+    plt.colorbar()
+    plt.subplot(4, 3, 11)
+    plt.title('cut slice xz')
+    plt.imshow(F_slice_cut_xz)
+    plt.colorbar()
+    plt.subplot(4, 3, 12)
+    plt.title('difference')
+    plt.imshow(F_slice_cut_xz-F_full_cut_xz)
+    plt.colorbar()
+    plt.tight_layout()
+
     plt.savefig("image.pdf", bbox_inches='tight')
 
 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))
+error_cut_xy = np.max(np.abs(F_slice_cut_xy-F_full_cut_xy)) / np.max(np.abs(F_full_cut_xy))
+error_cut_xz = np.max(np.abs(F_slice_cut_xz-F_full_cut_xz)) / np.max(np.abs(F_full_cut_xz))
 
 print("F_full.shape", F_full.shape)
 print("F_slice_xz.shape", F_slice_xz.shape)
 print("F_slice_yz.shape", F_slice_yz.shape)
+print("F_full_cut_xy.shape", F_full_cut_xy.shape)
+print("F_full_cut_xz.shape", F_full_cut_xz.shape)
 print("error_xz", error_xz)
 print("error_yz", error_yz)
+print("error_cut_xy", error_cut_xy)
+print("error_cut_xz", error_cut_xz)
 
 assert(error_xz < 3.e-14)
 assert(error_yz < 3.e-14)
+assert(error_cut_xy < 3.e-14)
+assert(error_cut_xz < 3.e-14)

src/diagnostics/Diagnostic.H

@@ -102,6 +102,12 @@ private:
     bool m_include_ghost_cells = false; /**< if ghost cells are included in output */
     bool m_initialized = false; /**< if this object is fully initialized */
     std::vector<bool> m_has_field; /**< if there is field output to write */
+    bool m_use_custom_size_lo = false; /**< if a user defined diagnostics size should be used (lo)*/
+    bool m_use_custom_size_hi = false; /**< if a user defined diagnostics size should be used (hi)*/
+    /** 3D array with lower ends of the diagnostics grid */
+    amrex::Array<amrex::Real, 3> m_diag_lo {0., 0., 0.};
+    /** 3D array with upper ends of the diagnostics grid */
+    amrex::Array<amrex::Real, 3> m_diag_hi {0., 0., 0.};
 };
 
 #endif // DIAGNOSTIC_H_

src/diagnostics/Diagnostic.cpp

@@ -33,6 +33,9 @@ Diagnostic::Diagnostic (int nlev)
 
     queryWithParser(ppd, "include_ghost_cells", m_include_ghost_cells);
 
+    m_use_custom_size_lo = queryWithParser(ppd, "patch_lo", m_diag_lo);
+    m_use_custom_size_hi = queryWithParser(ppd, "patch_hi", m_diag_hi);
+
     for(int ilev = 0; ilev<nlev; ++ilev) {
         amrex::Array<int,3> diag_coarsen_arr{1,1,1};
         // set all levels the same for now
@@ -128,6 +131,31 @@ Diagnostic::ResizeFDiagFAB (amrex::Box local_box, amrex::Box domain, const int l
         domain.grow(Fields::m_slices_nguards);
     }
 
+    {
+        // shrink box to user specified bounds m_diag_lo and m_diag_hi (in real space)
+        const amrex::Real poff_x = GetPosOffset(0, geom, geom.Domain());
+        const amrex::Real poff_y = GetPosOffset(1, geom, geom.Domain());
+        const amrex::Real poff_z = GetPosOffset(2, geom, geom.Domain());
+        amrex::Box cut_domain = domain;
+        if (m_use_custom_size_lo) {
+            cut_domain.setSmall({
+                static_cast<int>(std::round((m_diag_lo[0] - poff_x)/geom.CellSize(0))),
+                static_cast<int>(std::round((m_diag_lo[1] - poff_y)/geom.CellSize(1))),
+                static_cast<int>(std::round((m_diag_lo[2] - poff_z)/geom.CellSize(2)))
+            });
+        }
+        if (m_use_custom_size_hi) {
+            cut_domain.setBig({
+                static_cast<int>(std::round((m_diag_hi[0] - poff_x)/geom.CellSize(0))),
+                static_cast<int>(std::round((m_diag_hi[1] - poff_y)/geom.CellSize(1))),
+                static_cast<int>(std::round((m_diag_hi[2] - poff_z)/geom.CellSize(2)))
+            });
+        }
+        // calculate intersection of boxes to prevent them getting larger
+        domain &= cut_domain;
+        local_box &= domain;
+    }
+
     amrex::RealBox diag_domain = geom.ProbDomain();
     for(int dir=0; dir<=2; ++dir) {
         // make diag_domain correspond to box

src/diagnostics/OpenPMDWriter.H

@@ -108,9 +108,13 @@ private:
     /** openPMD backend: h5, bp, or json. Default depends on what is available */
     std::string m_openpmd_backend = "default";
 
-    /** Last iteration that was written to file.
+    /** Last iteration that had fields written to file.
      * This is stored to make sure we don't write the last iteration multiple times. */
-    amrex::Vector<int> m_last_output_dumped;
+    amrex::Vector<int> m_last_field_output_dumped;
+
+    /** Last iteration that had beams written to file.
+     * This is stored to make sure we don't write the last iteration multiple times. */
+    amrex::Vector<int> m_last_beam_output_dumped;
 
     /** vector of length nbeams with the numbers of particles already written to file */
     amrex::Vector<uint64_t> m_offset;

src/diagnostics/OpenPMDWriter.cpp

@@ -58,7 +58,8 @@ OpenPMDWriter::InitDiagnostics (const int output_step, const int output_period,
        (!(output_step == max_step) && output_step % output_period != 0)) return;
 
     m_outputSeries.resize(nlev);
-    m_last_output_dumped.resize(nlev);
+    m_last_field_output_dumped.resize(nlev);
+    m_last_beam_output_dumped.resize(nlev);
 
     for (int lev=0; lev<nlev; ++lev) {
         std::string filename = m_file_prefix +
@@ -67,7 +68,8 @@ OpenPMDWriter::InitDiagnostics (const int output_step, const int output_period,
 
         m_outputSeries[lev] = std::make_unique< openPMD::Series >(
             filename, openPMD::Access::CREATE);
-        m_last_output_dumped[lev] = -1;
+        m_last_field_output_dumped[lev] = -1;
+        m_last_beam_output_dumped[lev] = -1;
     }
 
     // TODO: meta-data: author, mesh path, extensions, software
@@ -91,13 +93,13 @@ OpenPMDWriter::WriteDiagnostics (
             if (lev == 0) {
                 WriteBeamParticleData(a_multi_beam, iteration, output_step, it, a_box_sorter_vec,
                                       geom3D[lev], beamnames, lev);
+                m_last_beam_output_dumped[lev] = output_step;
             }
             m_outputSeries[lev]->flush();
-
         } else if (call_type == OpenPMDWriterCallType::fields && write_fields[lev]) {
             WriteFieldData(a_mf[lev], geom, slice_dir, varnames, iteration, output_step, lev);
             m_outputSeries[lev]->flush();
-            m_last_output_dumped[lev] = output_step;
+            m_last_field_output_dumped[lev] = output_step;
         }
     }
 }
@@ -151,18 +153,15 @@ OpenPMDWriter::WriteFieldData (
 
         // data type and global size of the simulation
         openPMD::Datatype datatype = openPMD::determineDatatype< amrex::Real >();
-        amrex::IntVect prob_size = data_box.bigEnd() - data_box.smallEnd()
-                                   + amrex::IntVect::TheUnitVector();
         amrex::Vector<std::uint64_t> probsize_reformat = {AMREX_D_DECL(
                      static_cast<std::uint64_t>(geom[lev].Domain().size()[2]),
-                     static_cast<std::uint64_t>(prob_size[1]),
-                     static_cast<std::uint64_t>(prob_size[0])
-                      )};
-        openPMD::Extent global_size = probsize_reformat; //geom[lev].Domain().size());
+                     static_cast<std::uint64_t>(geom[lev].Domain().size()[1]),
+                     static_cast<std::uint64_t>(geom[lev].Domain().size()[0]))};
+        openPMD::Extent global_size = probsize_reformat;
         // If slicing requested, remove number of points for the slicing direction
         if (slice_dir >= 0) global_size.erase(global_size.begin() + 2-slice_dir);
 
-        if (m_last_output_dumped[lev] != output_step) {
+        if (m_last_field_output_dumped[lev] != output_step) {
             openPMD::Dataset dataset(datatype, global_size);
             field_comp.resetDataset(dataset);
         }
@@ -172,7 +171,8 @@ OpenPMDWriter::WriteFieldData (
         data = openPMD::shareRaw( fab.dataPtr( icomp ) ); // non-owning view until flush()
 
         // Determine the offset and size of this data chunk in the global output
-        amrex::IntVect const box_offset = {0, 0, data_box.smallEnd(2)};
+        amrex::IntVect const box_offset =
+            {0, 0, data_box.smallEnd(2) - geom[lev].Domain().smallEnd(2)};
         openPMD::Offset chunk_offset = utils::getReversedVec(box_offset);
         openPMD::Extent chunk_size = utils::getReversedVec(data_box.size());
         if (slice_dir >= 0) { // remove Ny components
@@ -206,7 +206,7 @@ OpenPMDWriter::WriteBeamParticleData (MultiBeam& beams, openPMD::Iteration itera
         auto& beam = beams.getBeam(ibeam);
 
         const unsigned long long np = beams.get_total_num_particles(ibeam);
-        if (m_last_output_dumped[lev] != output_step) {
+        if (m_last_beam_output_dumped[lev] != output_step) {
             SetupPos(beam_species, beam, np, geom);
             SetupRealProperties(beam_species, m_real_names, np);
         }
@@ -408,8 +408,10 @@ OpenPMDWriter::SaveRealProperty (BeamParticleContainer& pc,
 void OpenPMDWriter::reset (const int output_step)
 {
     for (int lev = 0; lev<m_outputSeries.size(); ++lev) {
-        if (output_step != m_last_output_dumped[lev]) continue;
-        m_outputSeries[lev].reset();
+        if (output_step == m_last_field_output_dumped[lev] ||
+            output_step == m_last_beam_output_dumped[lev]) {
+            m_outputSeries[lev].reset();
+        }
     }
 }
 

tests/slice_IO.1Rank.sh

@@ -40,5 +40,21 @@ mpiexec -n 1 $HIPACE_EXECUTABLE $HIPACE_EXAMPLE_DIR/inputs_normalized \
         amr.n_cell = 64 88 100 \
         hipace.file_prefix=slice_io_yz
 
+mpiexec -n 1 $HIPACE_EXECUTABLE $HIPACE_EXAMPLE_DIR/inputs_normalized \
+        plasmas.sort_bin_size = 8 \
+        diagnostic.diag_type=xyz \
+        amr.n_cell = 64 88 100 \
+        diagnostic.patch_lo = -3 -100  0 \
+        diagnostic.patch_hi =  3  100  0 \
+        hipace.file_prefix=slice_io_cut_xy
+
+mpiexec -n 1 $HIPACE_EXECUTABLE $HIPACE_EXAMPLE_DIR/inputs_normalized \
+        plasmas.sort_bin_size = 8 \
+        diagnostic.diag_type=xz \
+        amr.n_cell = 64 88 100 \
+        diagnostic.patch_lo =  0 -3 -10 \
+        diagnostic.patch_hi =  4  3  10 \
+        hipace.file_prefix=slice_io_cut_xz
+
 # assert whether the two IO types match
 $HIPACE_EXAMPLE_DIR/analysis_slice_IO.py