Improvements on the adaptive time step calculation (#884)

* cleaning, first test looks ok

* fix compilation issue on GPU and add input parameter for tolerance

* predict next time step for better behavior in parallel simulations

* bit of cleaning

* further clarification

* cleaning

* add doc and comments and minor cleaning

* fix 1 compilation issue on GPU

* update benchmark

* fix benchmarks

* fix doc

* use the appropriate shape factor

* not in a kernel, let's use std::min/max

* updates from review

* improve init and adds flexibility

* proper init and remove flexibility in phase advance

* further init cleanup

* fix checksum

docs/source/run/parameters.rst

@@ -77,6 +77,27 @@ General parameters
     :math:`\omega_p` the plasma angular frequency and :math:`\gamma` is an average of Lorentz
     factors of the slowest particles in all beams.
 
+* ``hipace.adaptive_predict_step`` (`bool`) optional (default `0`)
+    Only used when using adaptive time step (see ``hipace.dt`` above).
+    If true, the current Lorentz factor and accelerating field on the beams are used to predict the (adaptive) ``dt`` of the next time steps.
+    This prediction is used to better estimate the betatron frequency at the beginning of the next step performed by the current rank.
+    It improves accuracy for parallel simulations (with significant deceleration and/or z-dependent plasma profile).
+    Note: should be on by default once good defaults are determined.
+
+* ``hipace.adaptive_control_phase_advance`` (`bool`) optional (default `0`)
+    Only used when using adaptive time step (see ``hipace.dt`` above).
+    If true, a test on the phase advance sets the time step so it matches the phase advance expected for a uniform plasma (to a certain tolerance).
+    This should improve the accuracy in the presence of density gradients.
+    Note: should be on by default once good defaults are determined.
+
+* ``hipace.adaptive_phase_tolerance`` (`Real`) optional (default `5.e-4`)
+    Only used when using adaptive time step (see ``hipace.dt`` above) and ``adaptive_control_phase_advance``.
+    Tolerance for the controlled phase advance described above (lower is more accurate, but should result in more time steps).
+
+* ``hipace.adaptive_phase_substeps`` (`int`) optional (default `100`)
+    Only used when using adaptive time step (see ``hipace.dt`` above) and ``adaptive_control_phase_advance``.
+    Number of sub-steps in the controlled phase advance described above (higher is more accurate, but should be slower).
+
 * ``hipace.normalized_units`` (`bool`) optional (default `0`)
     Using normalized units in the simulation.
 

src/Hipace.cpp

@@ -316,13 +316,23 @@ Hipace::InitData ()
     constexpr int lev = 0;
     m_initial_time = m_multi_beam.InitData(geom[lev]);
     m_multi_plasma.InitData(m_slice_ba, m_slice_dm, m_slice_geom, geom);
-    m_adaptive_time_step.Calculate(m_dt, m_multi_beam, m_multi_plasma.maxDensity());
+    if (Hipace::HeadRank()) {
+        m_adaptive_time_step.Calculate(m_physical_time, m_dt, m_multi_beam,
+                                       m_multi_plasma, geom[lev], m_fields);
+        m_adaptive_time_step.CalculateFromDensity(m_physical_time, m_dt, m_multi_plasma);
+    }
 #ifdef AMREX_USE_MPI
-    m_adaptive_time_step.WaitTimeStep(m_dt, m_comm_z);
-    m_adaptive_time_step.NotifyTimeStep(m_dt, m_comm_z);
+    m_adaptive_time_step.BroadcastTimeStep(m_dt, m_comm_z, m_numprocs_z);
 #endif
     m_physical_time = m_initial_time;
-
+#ifdef AMREX_USE_MPI
+    m_physical_time += m_dt * (m_numprocs_z-1-amrex::ParallelDescriptor::MyProc());
+#endif
+    if (!Hipace::HeadRank()) {
+        m_adaptive_time_step.Calculate(m_physical_time, m_dt, m_multi_beam,
+                                       m_multi_plasma, geom[lev], m_fields);
+        m_adaptive_time_step.CalculateFromDensity(m_physical_time, m_dt, m_multi_plasma);
+    }
     m_fields.checkInit();
 }
 
@@ -477,6 +487,8 @@ Hipace::Evolve ()
                 Notify(step, it); // just send signal to finish simulation
                 if (m_physical_time > m_max_time) break;
             }
+            m_adaptive_time_step.CalculateFromDensity(m_physical_time, m_dt, m_multi_plasma);
+
             // adjust time step to reach max_time
             m_dt = std::min(m_dt, m_max_time - m_physical_time);
 
@@ -524,8 +536,9 @@ Hipace::Evolve ()
             m_multi_beam.RemoveGhosts();
 
             if (m_physical_time < m_max_time) {
-                m_adaptive_time_step.Calculate(m_dt, m_multi_beam, m_multi_plasma.maxDensity(),
-                                               it, m_box_sorters, false);
+                m_adaptive_time_step.Calculate(
+                    m_physical_time, m_dt, m_multi_beam, m_multi_plasma,
+                    geom[lev], m_fields, it, m_box_sorters, false);
             } else {
                 m_dt = 2.*m_max_time;
             }

src/particles/particles_utils/FieldGather.H

@@ -338,4 +338,27 @@ void doLaserGatherShapeN (const amrex::Real xp,
     }
 }
 
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+void doGatherEz (const amrex::Real xp,
+                     const amrex::Real yp,
+                     amrex::Real& Ezp,
+                     Array3<amrex::Real const> const& slice_arr,
+                     const int ez_comp,
+                     const amrex::Real dx_inv,
+                     const amrex::Real dy_inv,
+                     const amrex::Real x_pos_offset,
+                     const amrex::Real y_pos_offset)
+{
+    // x,y direction
+    const amrex::Real x = (xp - x_pos_offset) * dx_inv;
+    const amrex::Real y = (yp - y_pos_offset) * dy_inv;
+
+    // Compute shape factors
+    auto [shape_y, j] = compute_single_shape_factor<false, 0>(y, 0);
+    auto [shape_x, i] = compute_single_shape_factor<false, 0>(x, 0);
+
+    // Gather Ez on particle from field on grid
+    Ezp += (shape_x * shape_y) * slice_arr(i, j, ez_comp);
+}
+
 #endif // FIELDGATHER_H_

src/particles/plasma/MultiPlasma.H

@@ -71,7 +71,7 @@ public:
      * the max is taken across species AND include m_adaptive_density, giving a way to
      * specify a density to the adaptive time step calculator even with no plasma species.
      */
-    amrex::Real maxDensity () const;
+    amrex::Real maxDensity (amrex::Real z) const;
 
     /** \brief Gather field values and push particles
      *

src/particles/plasma/MultiPlasma.cpp

@@ -70,12 +70,11 @@ MultiPlasma::InitData (amrex::Vector<amrex::BoxArray> slice_ba,
 }
 
 amrex::Real
-MultiPlasma::maxDensity () const
+MultiPlasma::maxDensity (amrex::Real z) const
 {
     amrex::Real max_density = 0;
-    const amrex::Real c_t = get_phys_const().c * Hipace::m_physical_time;
     for (auto& plasma : m_all_plasmas) {
-        max_density = amrex::max<amrex::Real>(max_density, plasma.m_density_func(0., 0., c_t));
+        max_density = amrex::max<amrex::Real>(max_density, plasma.m_density_func(0., 0., z));
     }
     return amrex::max(max_density, m_adaptive_density);
 }

src/utils/AdaptiveTimeStep.H

@@ -9,6 +9,7 @@
 #define ADAPTIVETIMESTEP_H_
 
 #include "particles/beam/MultiBeam.H"
+#include "particles/plasma/MultiPlasma.H"
 #include "particles/beam/BeamParticleContainer.H"
 #include <AMReX_AmrCore.H>
 
@@ -22,7 +23,21 @@ private:
 
     /** Number of time steps per betatron period for the adaptive time step */
     amrex::Real m_nt_per_betatron = 40.;
+    /** Upper bound of the time step. Avoid gigantic time step(s) when beam starts near vacuum */
     amrex::Real m_dt_max = std::numeric_limits<amrex::Real>::infinity();
+    /** uz of the slowest particles */
+    amrex::Real m_min_uz = std::numeric_limits<amrex::Real>::max();
+    /** Whether to predict the next time steps. More accurate for parallel simulations */
+    bool m_adaptive_predict_step = false;
+    /** If true, a test on the phase advance sets the time step so it matches the phase advance
+     * expected for a uniform plasma. Relevant in the presence of density gradients.
+     * The tolerance on the phase advance difference is controlled by m_adaptive_phase_tolerance */
+    bool m_adaptive_control_phase_advance = false;
+    /** Phase shift tolerance. Relevant when density gradients are present.
+     * Lower is more accurate. */
+    amrex::Real m_adaptive_phase_tolerance = 5.e-4;
+    /** Number of substeps on which the phase advance is monitored. */
+    int m_adaptive_phase_substeps = 100;
 
 public:
     /** Whether to use an adaptive time step */
@@ -31,31 +46,38 @@ public:
     explicit AdaptiveTimeStep (const int nbeams);
 
 #ifdef AMREX_USE_MPI
-    /** sends the calculated initial time step to the rank downstream
-     * \param[in] dt initial time step
-     * \param[in] a_comm_z longitudinal MPI communicator
-     */
-    void NotifyTimeStep (amrex::Real dt, MPI_Comm a_comm_z);
-
-    /** receives the calculated initial time step to thefrom the rank upstream
+    /** Head rank initial time step
      * \param[in,out] dt initial time step
      * \param[in] a_comm_z longitudinal MPI communicator
+     * \param[in] a_numprocs_z number of ranks in the z direction
      */
-    void WaitTimeStep (amrex::Real& dt, MPI_Comm a_comm_z);
+    void BroadcastTimeStep (amrex::Real& dt, MPI_Comm a_comm_z, int a_numprocs_z);
 #endif
 
     /** calculate the adaptive time step based on the beam energy
+     * \param[in] t current physical time
      * \param[in,out] dt the time step
      * \param[in] beams multibeam containing all beams
-     * \param[in] plasma_density maximum plasma density
+     * \param[in] plasmas multiplasma to get density profile info
+     * \param[in] geom geometry object
+     * \param[in] fields field object
      * \param[in] it current box number
      * \param[in] a_box_sorter_vec Vector (over species) of particles sorted by box
      * \param[in] initial whether to calculate the initial dt
      */
     void
-    Calculate (amrex::Real& dt, MultiBeam& beams, amrex::Real plasma_density, const int it=0,
+    Calculate (amrex::Real t, amrex::Real& dt, MultiBeam& beams, MultiPlasma& plasmas,
+               const amrex::Geometry& geom, const Fields& fields, const int it=0,
                const amrex::Vector<BoxSorter>& a_box_sorter_vec={}, const bool initial=true);
 
+    /** Right before starting a time step, correct its dt to account for local plasma density
+     * and resolve density gradients.
+     * \param[in] t current physical time
+     * \param[in,out] dt the time step
+     * \param[in] plasmas multiplasma to get density profile info
+     */
+    void
+    CalculateFromDensity (amrex::Real t, amrex::Real& dt, MultiPlasma& plasmas);
 };
 
 #endif // ADAPTIVETIMESTEP_H_