Add option to increase `max_step` or `stop_time` to fill BTD (#3693)

* add option to increase max_step to fill BTD

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

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

* make sure code compiles locally

* allow for unspecified max step if setting by BTD

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

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

* clarify docs, match default max_step value

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

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

* Apply suggestions from code review

shortening flag name

* add example input script, fix bugf

* check stop_time and max_step

* account for stop time

* Fix logic, improve documentation

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

5 files changed, 49 insertions, 8 deletions

diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst
index 84990e698..5fae9c383 100644
--- a/Docs/source/usage/parameters.rst
+++ b/Docs/source/usage/parameters.rst
@@ -57,6 +57,15 @@ Overall simulation parameters
     printed to standard output. Currently only works if the Lorentz boost and
     the moving window are along the z direction.
 
+* ``warpx.compute_max_step_from_btd`` (`integer`; 0 by default) optional
+    Can be useful when computing back-transformed diagnostics.  If specified,
+    automatically calculates the number of iterations required in the boosted
+    frame for all back-transformed diagnostics to be completed. If ``max_step``,
+    ``stop_time``, or ``warpx.zmax_plasma_to_compute_max_step`` are not specified,
+    or the current values of ``max_step`` and/or ``stop_time`` are too low to fill
+    all BTD snapshots, the values of ``max_step`` and/or ``stop_time`` are
+    overwritten with the new values and printed to standard output.
+
 * ``warpx.random_seed`` (`string` or `int` > 0) optional
     If provided ``warpx.random_seed = random``, the random seed will be determined
     using `std::random_device` and `std::clock()`,
diff --git a/Examples/Tests/btd_rz/inputs_rz_z_boosted_BTD b/Examples/Tests/btd_rz/inputs_rz_z_boosted_BTD
index a2bdc8089..ce2675ead 100644
--- a/Examples/Tests/btd_rz/inputs_rz_z_boosted_BTD
+++ b/Examples/Tests/btd_rz/inputs_rz_z_boosted_BTD
@@ -1,5 +1,5 @@
 # Maximum number of time steps
-warpx.zmax_plasma_to_compute_max_step = 500e-6
+warpx.compute_max_step_from_btd = 1
 # number of grid points
 amr.n_cell = 32 256
 
diff --git a/Source/Diagnostics/BTDiagnostics.cpp b/Source/Diagnostics/BTDiagnostics.cpp
index 84a0501b5..ab538efb5 100644
--- a/Source/Diagnostics/BTDiagnostics.cpp
+++ b/Source/Diagnostics/BTDiagnostics.cpp
@@ -42,7 +42,7 @@
 #include <cmath>
 #include <cstdio>
 #include <memory>
-#include <string>
+#include <sstream>
 #include <vector>
 
 using namespace amrex::literals;
@@ -158,15 +158,38 @@ void BTDiagnostics::DerivedInitData ()
     // j >= i / gamma / (1+beta) * dt_snapshot / dt_boosted_frame
     const int final_snapshot_starting_step = static_cast<int>(std::ceil(final_snapshot_iteration / warpx.gamma_boost / (1._rt+warpx.beta_boost) * m_dt_snapshots_lab / dt_boosted_frame));
     const int final_snapshot_fill_iteration = final_snapshot_starting_step + num_buffers * m_buffer_size - 1;
-    if (final_snapshot_fill_iteration > warpx.maxStep()) {
-        std::string warn_string =
-            "\nSimulation might not run long enough to fill all BTD snapshots.\n"
-            "Final step: " + std::to_string(warpx.maxStep()) + "\n"
-            "Last BTD snapshot fills around step: " + std::to_string(final_snapshot_fill_iteration);
+    const amrex::Real final_snapshot_fill_time = final_snapshot_fill_iteration * dt_boosted_frame;
+    if (warpx.compute_max_step_from_btd) {
+        if (final_snapshot_fill_iteration > warpx.maxStep()) {
+            warpx.updateMaxStep(final_snapshot_fill_iteration);
+            amrex::Print()<<"max_step insufficient to fill all BTD snapshots. Automatically increased to: "
+                << final_snapshot_fill_iteration << std::endl;
+
+        }
+        if (final_snapshot_fill_time > warpx.stopTime()) {
+            warpx.updateStopTime(final_snapshot_fill_time);
+            amrex::Print()<<"stop_time insufficient to fill all BTD snapshots. Automatically increased to: "
+                << final_snapshot_fill_time << std::endl;
+
+        }
+        if (warpx.maxStep() == std::numeric_limits<int>::max() && warpx.stopTime() == std::numeric_limits<amrex::Real>::max()) {
+            amrex::Print()<<"max_step unspecified and stop time unspecified.  Setting max step to "
+                <<final_snapshot_fill_iteration<< " to fill all BTD snapshots." << std::endl;
+            warpx.updateMaxStep(final_snapshot_fill_iteration);
+        }
+
+    } else if (final_snapshot_fill_iteration > warpx.maxStep() || final_snapshot_fill_time > warpx.stopTime()) {
+        std::stringstream warn_string;
+            warn_string << "\nSimulation might not run long enough to fill all BTD snapshots.\n"
+            << "Final step: " << warpx.maxStep() << "\n"
+            <<"Stop time: " << warpx.stopTime() << "\n"
+            <<"Last BTD snapshot fills around step: " << final_snapshot_fill_iteration << "\n"
+            <<" or time: " << final_snapshot_fill_time << "\n";
         ablastr::warn_manager::WMRecordWarning(
-            "BTD", warn_string,
+            "BTD", warn_string.str(),
             ablastr::warn_manager::WarnPriority::low);
     }
+
 #ifdef WARPX_DIM_RZ
     UpdateVarnamesForRZopenPMD();
 #endif
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 2b8d524d0..fe30c9b5c 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -295,6 +295,9 @@ public:
     //! simulation domain along z reaches #zmax_plasma_to_compute_max_step in the boosted frame
     static bool do_compute_max_step_from_zmax;
 
+    //! If true, the code will compute max_step from the back transformed diagnostics
+    static bool compute_max_step_from_btd;
+
     static bool do_dynamic_scheduling;
     static bool refine_plasma;
 
@@ -779,7 +782,9 @@ public:
     bool getis_synchronized() const {return is_synchronized;}
 
     int maxStep () const {return max_step;}
+    void updateMaxStep (const int new_max_step) {max_step = new_max_step;}
     amrex::Real stopTime () const {return stop_time;}
+    void updateStopTime (const amrex::Real new_stop_time) {stop_time = new_stop_time;}
 
     void AverageAndPackFields( amrex::Vector<std::string>& varnames,
         amrex::Vector<amrex::MultiFab>& mf_avg, const amrex::IntVect ngrow) const;
diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index 48037eaa1..cf5892f15 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -114,6 +114,7 @@ Real WarpX::gamma_boost = 1._rt;
 Real WarpX::beta_boost = 0._rt;
 Vector<int> WarpX::boost_direction = {0,0,0};
 bool WarpX::do_compute_max_step_from_zmax = false;
+bool WarpX::compute_max_step_from_btd = false;
 Real WarpX::zmax_plasma_to_compute_max_step = 0._rt;
 
 short WarpX::current_deposition_algo;
@@ -622,6 +623,9 @@ WarpX::ReadParameters ()
             pp_warpx, "zmax_plasma_to_compute_max_step",
             zmax_plasma_to_compute_max_step);
 
+        pp_warpx.query("compute_max_step_from_btd",
+            compute_max_step_from_btd);
+
         pp_warpx.query("do_moving_window", do_moving_window);
         if (do_moving_window)
         {