[Draft] Transport and Covariance Matrices

src/initialization/InitDistribution.cpp

@@ -10,6 +10,7 @@
 #include "initialization/InitDistribution.H"
 
 #include "ImpactX.H"
+#include "particles/CovarianceMatrix.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/distribution/All.H"
 
@@ -32,6 +33,48 @@
 namespace impactx
 {
 
+    /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix
+     */
+    CovarianceMatrix
+    create_covariance_matrix (
+        distribution::KnownDistributions const & distr
+    )
+    {
+        // zero out the 6x6 matrix
+        CovarianceMatrix cv;
+        for (int i=1; i<=6; ++i) {
+            for (int j = 1; j <= 6; ++j) {
+                cv(i, j) = 0.0;
+            }
+        }
+
+        // initialize from 2nd order beam moments
+        std::visit([&](auto&& distribution) {
+            // quick hack
+            using Distribution = std::remove_cv_t< std::remove_reference_t< decltype(distribution)> >;
+            if constexpr (std::is_same<Distribution, distribution::Empty>::value ||
+                          std::is_same<Distribution, distribution::Thermal>::value)
+            {
+                throw std::runtime_error("Empty and Thermal type cannot create Covariance matrices!");
+            } else {
+                amrex::ParticleReal lambdaX = distribution.m_lambdaX;
+                amrex::ParticleReal lambdaY = distribution.m_lambdaY;
+                amrex::ParticleReal lambdaT = distribution.m_lambdaT;
+                amrex::ParticleReal lambdaPx;
+                amrex::ParticleReal lambdaPy;
+                amrex::ParticleReal lambdaPt;
+                amrex::ParticleReal muxpx;
+                amrex::ParticleReal muypy;
+                amrex::ParticleReal mutpt;
+
+                // convert to co-variance matrix
+                // TODO
+            }
+        }, distr);
+
+        return cv;
+    }
+
     void
     ImpactX::add_particles (
         amrex::ParticleReal bunch_charge,
@@ -95,7 +138,7 @@
             amrex::ParticleReal * const AMREX_RESTRICT py_ptr = py.data();
             amrex::ParticleReal * const AMREX_RESTRICT pt_ptr = pt.data();
 
-            using Distribution = std::remove_reference_t< std::remove_cv_t<decltype(distribution)> >;
+            using Distribution = std::remove_reference_t< std::remove_cv_t<decltype(distribution)> >; // TODO: switch order ov remove_ ...?
             initialization::InitSingleParticleData<Distribution> const init_single_particle_data(
                 distribution, x_ptr, y_ptr, t_ptr, px_ptr, py_ptr, pt_ptr);
             amrex::ParallelForRNG(npart_this_proc, init_single_particle_data);

src/initialization/InitElement.cpp

@@ -9,6 +9,7 @@
  */
 #include "ImpactX.H"
 #include "particles/elements/All.H"
+#include "particles/elements/mixin/lineartransport.H"
 
 #include <AMReX.H>
 #include <AMReX_BLProfiler.H>
@@ -434,6 +435,25 @@ namespace detail
                     read_element(sub_element_name, m_lattice, nslice_default, mapsteps_default);
                 }
             }
+        } else if (element_type == "linear_map")
+        {
+            // Parse the lattice elements for the sub-lattice in the line
+            amrex::ParmParse pp_sub_lattice(element_name);
+
+            auto a = detail::query_alignment(pp_element);
+
+            elements::LinearTransport::Map6x6 transport_map;
+            for (int i=1; i<=6; ++i) {
+                for (int j=1; j<=6; ++j) {
+                    amrex::ParticleReal R_ij = (i == j) ? 1.0 : 0.0;
+                    std::string name = "R" + std::to_string(i) + std::to_string(j);
+                    pp_element.queryAdd(name.c_str(), R_ij);
+
+                    transport_map(i, j) = R_ij;
+                }
+            }
+
+            m_lattice.emplace_back(LinearMap(transport_map, a["dx"], a["dy"], a["rotation_degree"]) );
         } else {
             amrex::Abort("Unknown type for lattice element " + element_name + ": " + element_type);
         }

src/particles/CovarianceMatrix.H

@@ -0,0 +1,24 @@
+/* Copyright 2022-2024 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H
+#define IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H
+
+#include <AMReX_Array.H>
+#include <AMReX_REAL.H>
+
+
+namespace impactx
+{
+    /** this is a 6x6 matrix */
+    using CovarianceMatrix = amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6>;
+
+} // namespace impactx::distribution
+
+#endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H

src/particles/PushAll.H

@@ -49,6 +49,10 @@ namespace impactx
             element(ref_part);
         }
 
+        // push covariance matrix
+        // TODO
+        // note: decide what to do for elements that have no covariance matrix
+
         // loop over refinement levels
         int const nLevel = pc.finestLevel();
         for (int lev = 0; lev <= nLevel; ++lev)

src/particles/distribution/Gaussian.H

@@ -134,7 +134,6 @@ namespace impactx::distribution
             pt = a2;
         }
 
-    private:
         amrex::ParticleReal m_lambdaX, m_lambdaY, m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx, m_lambdaPy, m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx, m_muypy, m_mutpt;  //! correlation length-momentum

src/particles/distribution/KVdist.H

@@ -147,7 +147,6 @@ namespace impactx::distribution
             pt = a2;
         }
 
-    private:
         amrex::ParticleReal m_lambdaX, m_lambdaY, m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx, m_lambdaPy, m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx, m_muypy, m_mutpt;  //! correlation length-momentum

src/particles/distribution/Kurth4D.H

@@ -157,7 +157,6 @@ namespace impactx::distribution
             pt = a2;
         }
 
-    private:
         amrex::ParticleReal m_lambdaX, m_lambdaY, m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx, m_lambdaPy, m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx, m_muypy, m_mutpt;  //! correlation length-momentum

src/particles/distribution/Kurth6D.H

@@ -163,7 +163,6 @@ namespace impactx::distribution
             pt = a2;
         }
 
-    private:
         amrex::ParticleReal m_lambdaX, m_lambdaY, m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx, m_lambdaPy, m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx, m_muypy, m_mutpt;  //! correlation length-momentum

src/particles/distribution/Semigaussian.H

@@ -146,7 +146,6 @@ namespace impactx::distribution
             pt = a2;
         }
 
-    private:
         amrex::ParticleReal m_lambdaX, m_lambdaY, m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx, m_lambdaPy, m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx, m_muypy, m_mutpt;  //! correlation length-momentum

src/particles/distribution/Triangle.H

@@ -153,7 +153,7 @@ namespace impactx::distribution
             t = a1;
             pt = a2;
         }
-    private:
+
         amrex::ParticleReal m_lambdaX, m_lambdaY, m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx, m_lambdaPy, m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx, m_muypy, m_mutpt; //! correlation length-momentum

src/particles/distribution/Waterbag.H

@@ -151,7 +151,6 @@ namespace impactx::distribution
             pt = a2;
         }
 
-    private:
         amrex::ParticleReal m_lambdaX,m_lambdaY,m_lambdaT;  //! related position axis intercepts (length) of the phase space ellipse
         amrex::ParticleReal m_lambdaPx,m_lambdaPy,m_lambdaPt;  //! related momentum axis intercepts of the phase space ellipse
         amrex::ParticleReal m_muxpx,m_muypy,m_mutpt;  //! correlation length-momentum

src/particles/elements/All.H

@@ -20,19 +20,20 @@
 #include "ConstF.H"
 #include "DipEdge.H"
 #include "Drift.H"
+#include "Empty.H"
 #include "ExactDrift.H"
 #include "ExactSbend.H"
 #include "Kicker.H"
+#include "LinearMap.H"
 #include "Multipole.H"
-#include "Empty.H"
 #include "NonlinearLens.H"
 #include "Programmable.H"
+#include "PRot.H"
 #include "Quad.H"
 #include "RFCavity.H"
 #include "Sbend.H"
 #include "ShortRF.H"
 #include "Sol.H"
-#include "PRot.H"
 #include "SoftSol.H"
 #include "SoftQuad.H"
 #include "TaperedPL.H"
@@ -60,6 +61,7 @@ namespace impactx
         ExactDrift,
         ExactSbend,
         Kicker,
+        LinearMap,
         Multipole,
         NonlinearLens,
         Programmable,

src/particles/elements/LinearMap.H

@@ -0,0 +1,104 @@
+/* Copyright 2022-2024 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_ELEMENT_LINEAR_MAP_H
+#define IMPACTX_ELEMENT_LINEAR_MAP_H
+
+#include "particles/ImpactXParticleContainer.H"
+#include "mixin/alignment.H"
+#include "mixin/beamoptic.H"
+#include "mixin/lineartransport.H"
+#include "mixin/thin.H"
+#include "mixin/nofinalize.H"
+
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+
+#include <cmath>
+
+
+namespace impactx
+{
+    struct LinearMap
+    : public elements::BeamOptic<LinearMap>,
+      public elements::Thin,
+      public elements::Alignment,
+      public elements::LinearTransport,
+      public elements::NoFinalize
+    {
+        static constexpr auto type = "LinearMap";
+        using PType = ImpactXParticleContainer::ParticleType;
+
+        /** A thin element that applies ...
+         *
+         * @param R user-provided transport map
+         * @param dx horizontal translation error in m
+         * @param dy vertical translation error in m
+         * @param rotation_degree rotation error in the transverse plane [degrees]
+         */
+        LinearMap (
+            LinearTransport::Map6x6 const & R,
+            amrex::ParticleReal dx = 0,
+            amrex::ParticleReal dy = 0,
+            amrex::ParticleReal rotation_degree = 0
+        )
+        : Alignment(dx, dy, rotation_degree)
+        {
+            for (int i=1; i<=6; ++i) {
+                for (int j = 1; j <= 6; ++j) {
+                    m_transport_map(i, j) = R(i, j);
+                }
+            }
+        }
+
+        /** Push all particles */
+        using BeamOptic::operator();
+
+        /** This is a ... functor, so that a variable of this type can be used like a
+         *  function.
+         *
+         * @param x particle position in x
+         * @param y particle position in y
+         * @param t particle position in t (unused)
+         * @param px particle momentum in x
+         * @param py particle momentum in y
+         * @param pt particle momentum in t (unused)
+         * @param idcpu particle global index
+         * @param refpart reference particle (unused)
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+            amrex::ParticleReal & AMREX_RESTRICT x,
+            amrex::ParticleReal & AMREX_RESTRICT y,
+            amrex::ParticleReal & AMREX_RESTRICT t,
+            amrex::ParticleReal & AMREX_RESTRICT px,
+            amrex::ParticleReal & AMREX_RESTRICT py,
+            amrex::ParticleReal & AMREX_RESTRICT pt,
+            [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu,
+            [[maybe_unused]] RefPart const & refpart
+        ) const
+        {
+            using namespace amrex::literals; // for _rt and _prt
+
+            // shift due to alignment errors of the element
+            shift_in(x, y, px, py);
+
+            // ... TODO ...
+
+            // undo shift due to alignment errors of the element
+            shift_out(x, y, px, py);
+        }
+
+        /** This pushes the reference particle. */
+        using Thin::operator();
+    };
+
+} // namespace impactx
+
+#endif // IMPACTX_ELEMENT_LINEAR_MAP_H

src/particles/elements/mixin/lineartransport.H

@@ -0,0 +1,50 @@
+/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H
+#define IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <ablastr/constant.H>
+
+#include <AMReX_Math.H>
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+
+
+namespace impactx::elements
+{
+    /** This is a helper class for lattice elements that can be expressed as linear transport maps.
+     */
+    struct LinearTransport
+    {
+        /** ...
+         */
+        LinearTransport (
+        )
+        {
+        }
+
+        //LinearTransport () = default;
+        LinearTransport (LinearTransport const &) = default;
+        LinearTransport& operator= (LinearTransport const &) = default;
+        LinearTransport (LinearTransport&&) = default;
+        LinearTransport& operator= (LinearTransport&& rhs) = default;
+
+        ~LinearTransport () = default;
+
+        // 6x6 linear transport map
+        using Map6x6 = amrex::Array2D<amrex::ParticleReal, 1, 6, 1, 6>;
+        Map6x6 m_transport_map; ///< linearized map
+    };
+
+} // namespace impactx::elements
+
+#endif // IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H