[Draft] track_covariance_map

src/ImpactX.H

@@ -104,6 +104,10 @@ namespace impactx
          */
         void track_particles ();
 
+        /** Run the linear transport of the covariance matrix simulation loop
+         */
+        void track_covariance_map ();
+
         /** Query input for warning logger variables and set up warning logger accordingly
          *
          * Input variables are: ``always_warn_immediately`` and ``abort_on_warning_threshold``.

src/ImpactX.cpp

@@ -336,4 +336,37 @@ namespace impactx {
             }, element_variant);
         }
     }
+
+    void
+    ImpactX::track_covariance_map ()
+    {
+        // TODO: move whole body out in separate file
+
+        // TODO: init done?
+        amrex::ParmParse const pp_dist = amrex::ParmParse("beam");
+        auto ref = initialization::read_reference_particle(pp_dist);
+        auto dist = initialization::read_distribution(pp_dist);
+        auto cm = impactx::initialization::create_covariance_matrix(dist);
+
+        // TODO: output of init state?
+
+        // loop over all beamline elements
+        for (auto & element_variant : m_lattice)
+        {
+            std::visit([&ref, &cm](auto&& element){
+                // push reference particle in global coordinates
+                {
+                    BL_PROFILE("impactx::Push::RefPart");
+                    element(ref);
+                }
+
+                // push Covariance Matrix
+                cm = element.transport_map(ref) * cm * element.transport_map(ref).transpose();
+
+                // TODO: later on, we can add element slicing and space charge kicking in this loop
+            }, element_variant);
+        }
+
+        // TODO: output
+    }
 } // namespace impactx

src/initialization/InitElement.cpp

@@ -529,7 +529,7 @@ namespace detail
             amrex::ParticleReal ds = 0.0;
             pp_element.queryAdd("ds", ds);
 
-            elements::mixin::LinearTransport::Map6x6 transport_map = elements::mixin::LinearTransport::Map6x6::Identity();
+            Map6x6 transport_map = Map6x6::Identity();
 
             // safe to ParmParse inputs for reproducibility
             for (int i=1; i<=6; ++i) {

src/particles/CovarianceMatrix.H

@@ -17,7 +17,10 @@
 namespace impactx
 {
     /** this is a 6x6 matrix */
-    using CovarianceMatrix = amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1>;
+    using Map6x6 = amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1>;
+
+    /** the covariance matrix is 6x6 */
+    using CovarianceMatrix = Map6x6;
 
 } // namespace impactx::distribution
 

src/particles/PushAll.H

@@ -52,7 +52,7 @@ namespace impactx
         }
 
         // push covariance matrix
-        // TODO
+        // TODO ?
         // note: decide what to do for elements that have no covariance matrix
 
         // loop over refinement levels

src/particles/elements/Aperture.H

@@ -12,6 +12,7 @@
 
 
 #include "particles/ImpactXParticleContainer.H"
+#include "particles/CovarianceMatrix.H"
 #include "mixin/alignment.H"
 #include "mixin/beamoptic.H"
 #include "mixin/thin.H"
@@ -179,6 +180,19 @@ namespace impactx
         /** This pushes the reference particle. */
         using Thin::operator();
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // nothing to do -- TODO: IS that true?
+            return Map6x6::Identity();
+        }
+
         Shape m_shape; //! aperture type (rectangular, elliptical)
         Action m_action; //! action type (transmit, absorb)
         amrex::ParticleReal m_aperture_x; //! maximum horizontal coordinate

src/particles/elements/Buncher.H

@@ -115,6 +115,19 @@ namespace impactx
         /** This pushes the reference particle. */
         using Thin::operator();
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_V; //! normalized (max) RF voltage drop.
         amrex::ParticleReal m_k; //! RF wavenumber in 1/m.
     };

src/particles/elements/CFbend.H

@@ -236,6 +236,19 @@ namespace impactx
             refpart.s = s + slice_ds;
         }
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_rc; //! bend radius in m
         amrex::ParticleReal m_k;  //! quadrupole strength in m^(-2)
     };

src/particles/elements/ChrDrift.H

@@ -182,6 +182,19 @@ namespace impactx
             // advance integrated path length
             refpart.s = s + slice_ds;
         }
+
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
     };
 
 } // namespace impactx

src/particles/elements/ChrPlasmaLens.H

@@ -221,6 +221,19 @@ namespace impactx
             refpart.s = s + slice_ds;
         }
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_k; //! focusing strength in 1/m^2 (or T/m)
         int m_unit; //! unit specification for focusing strength
     };

src/particles/elements/ChrQuad.H

@@ -260,6 +260,19 @@ namespace impactx
             refpart.s = s + slice_ds;
         }
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_k; //! quadrupole strength in 1/m^2 (or T/m)
         int m_unit; //! unit specification for quad strength
     };

src/particles/elements/ChrUniformAcc.H

@@ -246,6 +246,19 @@ namespace impactx
             refpart.s = s + slice_ds;
         }
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_ez; //! electric field strength in 1/m
         amrex::ParticleReal m_bz; //! magnetic field strength in 1/m
     };

src/particles/elements/ConstF.H

@@ -179,6 +179,19 @@ namespace impactx
 
         }
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_kx; //! focusing x strength in 1/m
         amrex::ParticleReal m_ky; //! focusing y strength in 1/m
         amrex::ParticleReal m_kt; //! focusing t strength in 1/m

src/particles/elements/DipEdge.H

@@ -121,6 +121,19 @@ namespace impactx
         /** This pushes the reference particle. */
         using Thin::operator();
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_psi; //! pole face angle in rad
         amrex::ParticleReal m_rc; //! bend radius in m
         amrex::ParticleReal m_g; //! gap parameter in m

src/particles/elements/Drift.H

@@ -32,6 +32,7 @@ namespace impactx
     struct Drift
     : public elements::mixin::Named,
       public elements::mixin::BeamOptic<Drift>,
+      public elements::mixin::LinearTransport<Drift>,
       public elements::mixin::Thick,
       public elements::mixin::Alignment,
       public elements::mixin::PipeAperture,
@@ -178,7 +179,7 @@ namespace impactx
          * @returns 6x6 transport matrix
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-        elements::mixin::LinearTransport::Map6x6
+        Map6x6
         transport_map (RefPart & AMREX_RESTRICT refpart) const
         {
             using namespace amrex::literals; // for _rt and _prt
@@ -191,7 +192,7 @@ namespace impactx
             amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt;
 
             // assign linear map matrix elements
-            elements::mixin::LinearTransport::Map6x6 R = elements::mixin::LinearTransport::Map6x6::Identity();
+            Map6x6 R = Map6x6::Identity();
             R(1,2) = slice_ds;
             R(3,4) = slice_ds;
             R(5,6) = slice_ds / betgam2;

src/particles/elements/Empty.H

@@ -51,6 +51,29 @@ namespace impactx
             // nothing to do
         }
 
+        /** Dos nothing to 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 function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // nothing to do
+            return Map6x6::Identity();
+        }
+
         /** Does nothing to a particle.
          *
          * @param x particle position in x

src/particles/elements/ExactDrift.H

@@ -143,8 +143,8 @@ namespace impactx
          * @param[in,out] refpart reference particle
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-        void operator() (RefPart & AMREX_RESTRICT refpart) const {
-
+        void operator() (RefPart & AMREX_RESTRICT refpart) const
+        {
             using namespace amrex::literals; // for _rt and _prt
 
             // assign input reference particle values
@@ -174,6 +174,19 @@ namespace impactx
             refpart.s = s + slice_ds;
 
         }
+
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
     };
 
 } // namespace impactx

src/particles/elements/ExactSbend.H

@@ -234,6 +234,19 @@ namespace impactx
 
         }
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_phi; //! bend angle in radians
         amrex::ParticleReal m_B;  //! magnetic field in T
     };

src/particles/elements/Kicker.H

@@ -139,6 +139,19 @@ namespace impactx
         /** This pushes the reference particle. */
         using Thin::operator();
 
+        /** This function returns the linear transport map.
+         *
+         * @param[in] refpart reference particle
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        Map6x6
+        transport_map ([[maybe_unused]] RefPart & AMREX_RESTRICT refpart) const
+        {
+            // TODO
+            return Map6x6::Identity();
+        }
+
         amrex::ParticleReal m_xkick; //! horizontal kick strength
         amrex::ParticleReal m_ykick; //! vertical kick strength
         UnitSystem m_unit; //! Kicks are for 0 dimensionless, or for 1 in T-m."