Make Drift Fast

Author: ax3l

src/elements/Drift.H

@@ -72,7 +72,31 @@ namespace impactx::elements
         /** Push all particles */
         using BeamOptic::operator();
 
+        /** Compute and cache the constants for the push.
+         *
+         * In particular, used to pre-compute and cache variables that are
+         * independent of the individually tracked particle.
+         *
+         * @param refpart reference particle
+         */
+        void compute_constants (RefPart const & refpart)
+        {
+            using namespace amrex::literals; // for _rt and _prt
+
+            Alignment::compute_constants(refpart);
+
+            // length of the current slice
+            m_slice_ds = m_ds / nslice();
+
+            // find beta*gamma^2
+            m_betgam2 = std::pow(refpart.pt, 2) - 1.0_prt;
+
+            m_slice_bg = m_slice_ds / m_betgam2;
+        }
+
         /** This is a drift functor, so that a variable of this type can be used like a drift function.
+         *
+         * The @see compute_constants method must be called before pushing particles through this operator.
          *
          * @param x particle position in x
          * @param y particle position in y
@@ -81,7 +105,7 @@ namespace impactx::elements
          * @param py particle momentum in y
          * @param pt particle momentum in t
          * @param idcpu particle global index
-         * @param refpart reference particle
+         * @param refpart reference particle (unused)
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
         void operator() (
@@ -92,7 +116,7 @@ namespace impactx::elements
             amrex::ParticleReal & AMREX_RESTRICT py,
             amrex::ParticleReal & AMREX_RESTRICT pt,
             uint64_t & AMREX_RESTRICT idcpu,
-            RefPart const & refpart
+            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart
         ) const
         {
             using namespace amrex::literals; // for _rt and _prt
@@ -108,19 +132,12 @@ namespace impactx::elements
             amrex::ParticleReal pyout = py;
             amrex::ParticleReal ptout = pt;
 
-            // length of the current slice
-            amrex::ParticleReal const slice_ds = m_ds / nslice();
-
-            // access reference particle values to find beta*gamma^2
-            amrex::ParticleReal const pt_ref = refpart.pt;
-            amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt;
-
             // advance position and momentum (drift)
-            xout = x + slice_ds * px;
+            xout = x + m_slice_ds * px;
             // pxout = px;
-            yout = y + slice_ds * py;
+            yout = y + m_slice_ds * py;
             // pyout = py;
-            tout = t + (slice_ds/betgam2) * pt;
+            tout = t + m_slice_bg * pt;
             // ptout = pt;
 
             // assign updated values
@@ -202,6 +219,13 @@ namespace impactx::elements
 
             return R;
         }
+
+    private:
+        // constants that are independent of the individually tracked particle,
+        // see: compute_constants() to refresh
+        amrex::ParticleReal m_slice_ds;  //! m_ds / nslice();
+        amrex::ParticleReal m_betgam2;   //! beta*gamma^2
+        amrex::ParticleReal m_slice_bg;  //! m_slice_ds / m_betgam2
     };
 
 } // namespace impactx