Added operator to push reference particle in global coordinates to al…

…l elements (#149) * Added math to push the reference particle in global coordinates to all existing elements. * Elements: Pass by Reference * Pass the reference particle by reference, so we can manipulate it in elements. * Remove true "no-op" implementations with a comment. * Add Doxygen comments. * Push the Reference Particle Co-authored-by: Axel Huebl <[email protected]>
ECP-WarpX · Jun 27, 2022 · d2476c8 · d2476c8
1 parent 1421304
commit d2476c8
Show file tree

Hide file tree

Showing 11 changed files with 188 additions and 10 deletions.
diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H
@@ -144,8 +144,8 @@ namespace impactx
          *
          * @returns refpart
          */
-        RefPart
-        GetRefParticle () const;
+        RefPart &
+        GetRefParticle ();
 
         /** Get particle shape
          */

diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp
@@ -123,8 +123,8 @@ namespace impactx
         m_refpart = refpart;
     }
 
-    RefPart
-    ImpactXParticleContainer::GetRefParticle () const
+    RefPart &
+    ImpactXParticleContainer::GetRefParticle ()
     {
         return m_refpart;
     }

diff --git a/src/particles/Push.cpp b/src/particles/Push.cpp
@@ -126,18 +126,20 @@ namespace detail
                 // ...
 
                 // preparing to access reference particle data: RefPart
-                RefPart ref_part;
-                ref_part = pc.GetRefParticle();
+                RefPart & ref_part = pc.GetRefParticle();
 
                 // loop over all beamline elements
                 for (auto & element_variant : lattice) {
                     // here we just access the element by its respective type
-                    std::visit([=](auto&& element) {
+                    std::visit([=, &ref_part](auto&& element) {
+                        // push beam particles relative to reference particle
                         detail::PushSingleParticle<decltype(element)> const pushSingleParticle(
                             element, aos_ptr, part_px, part_py, part_pt, ref_part);
-
-                        // loop over particles in the box
+                        //   loop over beam particles in the box
                         amrex::ParallelFor(np, pushSingleParticle);
+
+                        // push reference particle in global coordinates
+                        element(ref_part);
                     }, element_variant);
                 }; // end loop over all beamline elements
             } // end loop over all particle boxes

diff --git a/src/particles/elements/ConstF.H b/src/particles/elements/ConstF.H
@@ -88,6 +88,37 @@ namespace impactx
 
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                RefPart & AMREX_RESTRICT refpart) const {
+
+            using namespace amrex::literals; // for _rt and _prt
+
+            // assign input reference particle values
+            amrex::ParticleReal const x = refpart.x;
+            amrex::ParticleReal const px = refpart.px;
+            amrex::ParticleReal const y = refpart.y;
+            amrex::ParticleReal const py = refpart.py;
+            amrex::ParticleReal const z = refpart.z;
+            amrex::ParticleReal const pz = refpart.pz;
+            amrex::ParticleReal const t = refpart.t;
+            amrex::ParticleReal const pt = refpart.pt;
+
+            // assign intermediate parameter
+            amrex::ParticleReal const s = m_ds/sqrt(pow(pt,2)-1.0_prt);
+
+            // advance position and momentum (straight element)
+            refpart.x = x + s*px;
+            refpart.y = y + s*py;
+            refpart.z = z + s*pz;
+            refpart.t = t - s*pt;
+
+        }
+
     private:
         amrex::ParticleReal m_ds; //! segment length in m
         amrex::ParticleReal m_kx; //! focusing x strength in 1/m

diff --git a/src/particles/elements/DipEdge.H b/src/particles/elements/DipEdge.H
@@ -84,6 +84,17 @@ namespace impactx
             py = py + R43*y;
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                [[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const {
+
+            // nothing to do: this is a zero-length element
+        }
+
     private:
         amrex::ParticleReal m_psi; //! pole face angle in rad
         amrex::ParticleReal m_rc; //! bend radius in m

diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H
@@ -56,7 +56,7 @@ namespace impactx
             amrex::ParticleReal const y = p.pos(1);
             amrex::ParticleReal const t = p.pos(2);
 
-            // intialize output values of momenta
+            // initialize output values of momenta
             amrex::ParticleReal pxout = px;
             amrex::ParticleReal pyout = py;
             amrex::ParticleReal ptout = pt;
@@ -77,6 +77,38 @@ namespace impactx
             px = pxout;
             py = pyout;
             pt = ptout;
+
+        }
+
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                RefPart & AMREX_RESTRICT refpart) const {
+
+            using namespace amrex::literals; // for _rt and _prt
+
+            // assign input reference particle values
+            amrex::ParticleReal const x = refpart.x;
+            amrex::ParticleReal const px = refpart.px;
+            amrex::ParticleReal const y = refpart.y;
+            amrex::ParticleReal const py = refpart.py;
+            amrex::ParticleReal const z = refpart.z;
+            amrex::ParticleReal const pz = refpart.pz;
+            amrex::ParticleReal const t = refpart.t;
+            amrex::ParticleReal const pt = refpart.pt;
+
+            // assign intermediate parameter
+            amrex::ParticleReal const s = m_ds/sqrt(pow(pt,2)-1.0_prt);
+
+            // advance position and momentum (drift)
+            refpart.x = x + s*px;
+            refpart.y = y + s*py;
+            refpart.z = z + s*pz;
+            refpart.t = t - s*pt;
+
         }
 
     private:

diff --git a/src/particles/elements/Multipole.H b/src/particles/elements/Multipole.H
@@ -108,6 +108,17 @@ namespace impactx
 
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                [[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const {
+
+            // nothing to do: this is a zero-length element
+        }
+
     private:
         int m_multipole; //! multipole index
         int m_mfactorial; //! factorial of multipole index

diff --git a/src/particles/elements/NonlinearLens.H b/src/particles/elements/NonlinearLens.H
@@ -117,6 +117,17 @@ namespace impactx
 
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                [[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const {
+
+            // nothing to do: this is a zero-length element
+        }
+
     private:
         amrex::ParticleReal m_knll; //! integrated strength of the nonlinear lens (m)
         amrex::ParticleReal m_cnll; //! distance of singularities from the origin (m)

diff --git a/src/particles/elements/Quad.H b/src/particles/elements/Quad.H
@@ -101,6 +101,37 @@ namespace impactx
 
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                RefPart & AMREX_RESTRICT refpart) const {
+
+            using namespace amrex::literals; // for _rt and _prt
+
+            // assign input reference particle values
+            amrex::ParticleReal const x = refpart.x;
+            amrex::ParticleReal const px = refpart.px;
+            amrex::ParticleReal const y = refpart.y;
+            amrex::ParticleReal const py = refpart.py;
+            amrex::ParticleReal const z = refpart.z;
+            amrex::ParticleReal const pz = refpart.pz;
+            amrex::ParticleReal const t = refpart.t;
+            amrex::ParticleReal const pt = refpart.pt;
+
+            // assign intermediate parameter
+            amrex::ParticleReal const s = m_ds/sqrt(pow(pt,2)-1.0_prt);
+
+            // advance position and momentum (straight element)
+            refpart.x = x + s*px;
+            refpart.y = y + s*py;
+            refpart.z = z + s*pz;
+            refpart.t = t - s*pt;
+
+        }
+
     private:
         amrex::ParticleReal m_ds; //! segment length in m
         amrex::ParticleReal m_k; //! quadrupole strength in 1/m

diff --git a/src/particles/elements/Sbend.H b/src/particles/elements/Sbend.H
@@ -91,6 +91,43 @@ namespace impactx
 
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                RefPart & AMREX_RESTRICT refpart) const {
+
+            using namespace amrex::literals; // for _rt and _prt
+
+            // assign input reference particle values
+            amrex::ParticleReal const x = refpart.x;
+            amrex::ParticleReal const px = refpart.px;
+            amrex::ParticleReal const y = refpart.y;
+            amrex::ParticleReal const py = refpart.py;
+            amrex::ParticleReal const z = refpart.z;
+            amrex::ParticleReal const pz = refpart.pz;
+            amrex::ParticleReal const t = refpart.t;
+            amrex::ParticleReal const pt = refpart.pt;
+
+            // assign intermediate parameter
+            amrex::ParticleReal const theta = m_ds/m_rc;
+            amrex::ParticleReal const B = sqrt(pow(pt,2)-1.0_prt)/m_rc;
+
+            // advance position and momentum (bend)
+            refpart.px = px*cos(theta) - pz*sin(theta);
+            refpart.py = py;
+            refpart.pz = pz*cos(theta) + px*sin(theta);
+            refpart.pt = pt;
+
+            refpart.x = x + (refpart.pz - pz)/B;
+            refpart.y = y + (theta/B)*py;
+            refpart.z = z - (refpart.px - px)/B;
+            refpart.t = t - (theta/B)*pt;
+
+        }
+
     private:
         amrex::ParticleReal m_ds; //! segment length in m
         amrex::ParticleReal m_rc; //! bend radius in m

diff --git a/src/particles/elements/ShortRF.H b/src/particles/elements/ShortRF.H
@@ -85,6 +85,18 @@ namespace impactx
 
         }
 
+        /** This pushes the reference particle.
+         *
+         * @param[in,out] refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                [[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const {
+
+            // nothing to do: this is a zero-length element
+
+        }
+
     private:
         amrex::ParticleReal m_V; //! normalized (max) RF voltage drop.
         amrex::ParticleReal m_k; //! RF wavenumber in 1/m.