BTD diagnostics specified by intervals (#3367)

* BTD diagnostics specified by intervals

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address CI errors

* refactor and test

* no duplicate or out-of-order snapshots

* update documentation and tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean for CI

* still fighting CI

* Update Source/Diagnostics/BTDiagnostics.cpp

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Apply suggestions from code review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* refactor includes

* add an explicit constexpr

* Apply suggestions from code review

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* describe algorithm for BTD intervals list

* revert to old description of num_snapshots_lab

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* preserve behavior of num_snapshots_lab

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

7 files changed, 185 insertions, 7 deletions

diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst
index 36c6dc3b2..f87795a77 100644
--- a/Docs/source/usage/parameters.rst
+++ b/Docs/source/usage/parameters.rst
@@ -2159,6 +2159,18 @@ BackTransformed Diagnostics (with support for Plotfile/openPMD output)
 * ``<diag_name>.num_snapshots_lab`` (`integer`)
     Only used when ``<diag_name>.diag_type`` is ``BackTransformed``.
     The number of lab-frame snapshots that will be written.
+    Only this option or ``intervals`` should be specified;
+    a run-time error occurs if the user attempts to set both ``num_snapshots_lab`` and ``intervals``.
+
+* ``<diag_name>.intervals`` (`string`)
+    Only used when ``<diag_name>.diag_type`` is ``BackTransformed``.
+    Using the `Intervals parser`_ syntax, this string defines the lab frame times at which data is dumped,
+    given as multiples of the step size ``dt_snapshots_lab`` or ``dz_snapshots_lab`` described below.
+    Example: ``btdiag1.intervals = 10:11,20:24:2`` and ``btdiag1.dt_snapshots_lab = 1.e-12``
+    indicate to dump at lab times ``1e-11``, ``1.1e-11``, ``2e-11``, ``2.2e-11``, and ``2.4e-11`` seconds.
+    Note that the stop interval, the second number in the slice, must always be specified.
+    Only this option or ``num_snapshots_lab`` should be specified;
+    a run-time error occurs if the user attempts to set both ``num_snapshots_lab`` and ``intervals``.
 
 * ``<diag_name>.dt_snapshots_lab`` (`float`, in seconds)
     Only used when ``<diag_name>.diag_type`` is ``BackTransformed``.
diff --git a/Examples/Modules/RigidInjection/inputs_2d_BoostedFrame b/Examples/Modules/RigidInjection/inputs_2d_BoostedFrame
index 73301448e..bd8ce220a 100644
--- a/Examples/Modules/RigidInjection/inputs_2d_BoostedFrame
+++ b/Examples/Modules/RigidInjection/inputs_2d_BoostedFrame
@@ -53,7 +53,7 @@ diagnostics.diags_names = diag1 diag2
 
 diag1.diag_type = BackTransformed
 diag1.do_back_transformed_fields = 1
-diag1.num_snapshots_lab = 2
+diag1.intervals = :1
 diag1.dt_snapshots_lab = 1.8679589331096515e-13
 diag1.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz rho
 diag1.format = plotfile
diff --git a/Examples/Modules/boosted_diags/inputs_3d_slice b/Examples/Modules/boosted_diags/inputs_3d_slice
index 445bc8a80..c79389919 100644
--- a/Examples/Modules/boosted_diags/inputs_3d_slice
+++ b/Examples/Modules/boosted_diags/inputs_3d_slice
@@ -126,7 +126,7 @@ diag1.write_species = 1
 
 diag2.diag_type = BackTransformed
 diag2.do_back_transformed_fields = 1
-diag2.num_snapshots_lab = 4
+diag2.intervals = 0:4:2, 1:3:2
 diag2.dz_snapshots_lab = 0.001
 diag2.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz rho
 diag2.format = openpmd
diff --git a/Source/Diagnostics/BTDiagnostics.H b/Source/Diagnostics/BTDiagnostics.H
index beb6c17ec..fcfc8f13b 100644
--- a/Source/Diagnostics/BTDiagnostics.H
+++ b/Source/Diagnostics/BTDiagnostics.H
@@ -10,6 +10,7 @@
 #include "Diagnostics.H"
 #include "Diagnostics/ComputeDiagFunctors/ComputeDiagFunctor.H"
 #include "Utils/WarpXConst.H"
+#include "Utils/IntervalsParser.H"
 
 #include <AMReX_Box.H>
 #include <AMReX_Geometry.H>
@@ -37,6 +38,8 @@ private:
     bool m_plot_raw_fields_guards = false;
     /** Read relevant parameters for BTD */
     void ReadParameters ();
+    /** Determines timesteps at which BTD diagnostics are written to file */
+    BTDIntervalsParser m_intervals;
     /** \brief Flush m_mf_output and particles to file.
      * Currently, a temporary customized output format for the buffer
      * data is implemented and called in this function.
diff --git a/Source/Diagnostics/BTDiagnostics.cpp b/Source/Diagnostics/BTDiagnostics.cpp
index 1bd200fb4..beb367309 100644
--- a/Source/Diagnostics/BTDiagnostics.cpp
+++ b/Source/Diagnostics/BTDiagnostics.cpp
@@ -161,8 +161,18 @@ BTDiagnostics::ReadParameters ()
     WARPX_ALWAYS_ASSERT_WITH_MESSAGE(m_do_back_transformed_fields, " fields must be turned on for the new back-transformed diagnostics");
     if (m_do_back_transformed_fields == false) m_varnames.clear();
 
-    getWithParser(pp_diag_name, "num_snapshots_lab", m_num_snapshots_lab);
-    m_num_buffers = m_num_snapshots_lab;
+
+    std::vector<std::string> intervals_string_vec = {"0"};
+    bool const num_snapshots_specified = queryWithParser(pp_diag_name, "num_snapshots_lab", m_num_snapshots_lab);
+    bool const intervals_specified = pp_diag_name.queryarr("intervals", intervals_string_vec);
+    if (num_snapshots_specified)
+    {
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(!intervals_specified,
+            "For back-transformed diagnostics, user should specify either num_snapshots_lab or intervals, not both");
+        intervals_string_vec = {":" + std::to_string(m_num_snapshots_lab-1)};
+    }
+    m_intervals = BTDIntervalsParser(intervals_string_vec);
+    m_num_buffers = m_intervals.NumSnapshots();
 
     // Read either dz_snapshots_lab or dt_snapshots_lab
     bool snapshot_interval_is_specified = false;
@@ -241,7 +251,7 @@ BTDiagnostics::InitializeBufferData ( int i_buffer , int lev)
     auto & warpx = WarpX::GetInstance();
     // Lab-frame time for the i^th snapshot
     amrex::Real zmax_0 = warpx.Geom(lev).ProbHi(m_moving_window_dir);
-    m_t_lab.at(i_buffer) = i_buffer * m_dt_snapshots_lab
+    m_t_lab.at(i_buffer) = m_intervals.GetBTDIteration(i_buffer) * m_dt_snapshots_lab
         + m_gamma_boost*m_beta_boost*zmax_0/PhysConst::c;
 
     // Define buffer domain in boosted frame at level, lev, with user-defined lo and hi
diff --git a/Source/Utils/IntervalsParser.H b/Source/Utils/IntervalsParser.H
index c91fe271a..06258b109 100644
--- a/Source/Utils/IntervalsParser.H
+++ b/Source/Utils/IntervalsParser.H
@@ -21,7 +21,7 @@ public:
     * (0 for the starting point, std::numeric_limits<int>::max() for the stopping point and 1 for
     * the period). For example SliceParser(":1000:") is equivalent to SliceParser("0:1000:1").
     */
-    SliceParser (const std::string& instr);
+    SliceParser (const std::string& instr, bool isBTD=false);
 
     /**
     * \brief A method that returns true if the input integer is contained in the slice. (e.g. if
@@ -66,7 +66,14 @@ public:
     */
     int getStop () const;
 
+    /**
+     * @brief A method that returns the number of integers contained by the slice.
+     *
+     */
+    int numContained() const;
+
 private:
+    bool m_isBTD = false;
     int m_start = 0;
     int m_stop = std::numeric_limits<int>::max();
     int m_period = 1;
@@ -148,4 +155,53 @@ private:
     bool m_activated = false;
 };
 
+/**
+ * \brief This class is a parser for multiple slices of the form x,y,z,... where x, y and z are
+ * slices of the form i:j:k, as defined in the SliceParser class. This class contains a vector of
+ * SliceParsers. The supported function set differs from the IntervalsParser
+ */
+class BTDIntervalsParser
+{
+public:
+    /**
+    * \brief Default constructor of the BTDIntervalsParser class.
+    */
+    BTDIntervalsParser () = default;
+
+    /**
+    * \brief Constructor of the BTDIntervalsParser class.
+    *
+    * @param[in] instr_vec an input vector string, which when concatenated is of the form
+    * "x,y,z,...". This will call the constructor of SliceParser using x, y and z as input
+    * arguments.
+    */
+    BTDIntervalsParser (const std::vector<std::string>& instr_vec);
+
+    /**
+     * @brief Return the total number of unique labframe snapshots
+     */
+    int NumSnapshots ();
+
+    /**
+     * @brief Return the iteration number stored at index i_buffer
+     *
+     * @param i_buffer buffer or iteration index, between 0 and NumSnapshots
+     */
+    int GetBTDIteration(int i_buffer);
+
+    /**
+    * \brief A method that returns true if any of the slices contained by the IntervalsParser
+    * has a strictly positive period.
+    */
+    bool isActivated () const;
+
+private:
+    std::vector<int> m_btd_iterations;
+    std::vector<SliceParser> m_slices;
+    std::vector<int> m_slice_starting_i_buffer;
+    int m_n_snapshots;
+    static constexpr char m_separator = ',';
+    bool m_activated = false;
+};
+
 #endif // WARPX_INTERVALSPARSER_H_
diff --git a/Source/Utils/IntervalsParser.cpp b/Source/Utils/IntervalsParser.cpp
index e8425421f..4da0142a0 100644
--- a/Source/Utils/IntervalsParser.cpp
+++ b/Source/Utils/IntervalsParser.cpp
@@ -7,15 +7,18 @@
 #include <algorithm>
 #include <memory>
 
-SliceParser::SliceParser (const std::string& instr)
+SliceParser::SliceParser (const std::string& instr, const bool isBTD)
 {
+    m_isBTD = isBTD;
     // split string and trim whitespaces
     auto insplit = WarpXUtilStr::split<std::vector<std::string>>(instr, m_separator, true);
 
     if(insplit.size() == 1){ // no colon in input string. The input is the period.
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(!m_isBTD, "must specify interval stop for BTD");
         m_period = parseStringtoInt(insplit[0], "interval period");}
     else if(insplit.size() == 2) // 1 colon in input string. The input is start:stop
     {
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(!m_isBTD || !insplit[1].empty(), "must specify interval stop for BTD");
         if (!insplit[0].empty()){
             m_start = parseStringtoInt(insplit[0], "interval start");}
         if (!insplit[1].empty()){
@@ -23,6 +26,7 @@ SliceParser::SliceParser (const std::string& instr)
     }
     else // 2 colons in input string. The input is start:stop:period
     {
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(!m_isBTD || !insplit[1].empty(), "must specify interval stop for BTD");
         WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
             insplit.size() == 3,
             instr + "' is not a valid syntax for a slice.");
@@ -64,6 +68,8 @@ int SliceParser::getStart () const {return m_start;}
 
 int SliceParser::getStop () const {return m_stop;}
 
+int SliceParser::numContained () const {return (m_stop - m_start) / m_period + 1;}
+
 IntervalsParser::IntervalsParser (const std::vector<std::string>& instr_vec)
 {
     std::string inconcatenated;
@@ -116,3 +122,94 @@ int IntervalsParser::localPeriod (const int n) const
 }
 
 bool IntervalsParser::isActivated () const {return m_activated;}
+
+BTDIntervalsParser::BTDIntervalsParser (const std::vector<std::string>& instr_vec)
+{
+    std::string inconcatenated;
+    for (const auto& instr_element : instr_vec) inconcatenated +=instr_element;
+
+    auto const insplit = WarpXUtilStr::split<std::vector<std::string>>(inconcatenated, std::string(1,m_separator));
+
+    // parse the Intervals string into Slices and store each slice in m_slices,
+    // in order of increasing Slice start value
+    for(const auto& inslc : insplit)
+    {
+        bool isBTD = true;
+        SliceParser temp_slice(inslc, isBTD);
+        if (m_slices.size() > 0)
+        {
+            // find the last index i_slice where
+            // the start value of m_slices[i_slice] is greater than temp_slices' start_value
+            int i_slice = 0;
+            while (temp_slice.getStart() > m_slices[i_slice].getStart() && i_slice < static_cast<int>(m_slices.size()))
+            {
+                i_slice++;
+            }
+            m_slices.insert(m_slices.begin() + i_slice, temp_slice);
+        }
+        else
+        {
+            m_slices.push_back(temp_slice);
+        }
+    }
+    // from the vector of slices, m_slices,
+    // create a vector of integers, m_btd_iterations, containing
+    // the iteration of every back-transformed snapshot that will be saved
+    // the iteration values in m_btd_iterations are
+    // 1. saved in increasing order
+    // 2. unique, i.e. no duplicate iterations are saved
+    for (const auto& temp_slice : m_slices)
+    {
+        const int start = temp_slice.getStart();
+        const int period = temp_slice.getPeriod();
+        int btd_iter_ind;
+        // for Slice temp_slice in m_slices,
+        // determine the index in m_btd_iterations where temp_slice's starting value goes
+        //
+        // Implementation note:
+        // assuming the user mostly lists slices in ascending order,
+        // start at the end of m_btd_iterations and search backward
+        if (m_btd_iterations.size() == 0)
+        {
+            btd_iter_ind = 0;
+        }
+        else
+        {
+            btd_iter_ind = m_btd_iterations.size() - 1;
+            while (start < m_btd_iterations[btd_iter_ind] and btd_iter_ind>0)
+            {
+                btd_iter_ind--;
+            }
+        }
+        // insert each iteration contained in temp_slice into m_btd_iterations
+        // adding them in increasing sorted order and not adding any iterations
+        // already contained in m_btd_iterations
+        for (int ii = start; ii <= temp_slice.getStop(); ii += period)
+        {
+            if (m_btd_iterations.size() > 0)
+            {
+                // find where iteration ii should go in m_btd_iterations
+                while (ii > m_btd_iterations[btd_iter_ind] && btd_iter_ind < static_cast<int>(m_btd_iterations.size()))
+                {
+                    btd_iter_ind++;
+                }
+                if (ii != m_btd_iterations[btd_iter_ind])
+                {
+                    m_btd_iterations.insert(m_btd_iterations.begin() + btd_iter_ind, ii);
+                }
+            } else
+            {
+                m_btd_iterations.push_back(ii);
+            }
+        }
+        if ((temp_slice.getPeriod() > 0) &&
+               (temp_slice.getStop() >= start)) m_activated = true;
+    }
+}
+
+int BTDIntervalsParser::NumSnapshots () { return m_btd_iterations.size(); }
+
+int BTDIntervalsParser::GetBTDIteration(int i_buffer)
+{
+    return m_btd_iterations[i_buffer];
+}