[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,65 @@
 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 = distribution.m_lambdaPx;
+                amrex::ParticleReal lambdaPy = distribution.m_lambdaPy;
+                amrex::ParticleReal lambdaPt = distribution.m_lambdaPt;
+                amrex::ParticleReal muxpx = distribution.m_muxpx;
+                amrex::ParticleReal muypy = distribution.m_muypy;
+                amrex::ParticleReal mutpt = distribution.m_mutpt;
+
+                // use distribution inputs to populate a 6x6 covariance matrix
+                amrex::ParticleReal denom_x = 1.0 - muxpx*muxpx;
+                cv(1,1) = lambdaX*lambdaX / denom_x;
+                cv(1,2) = -lambdaX*lambdaPx*muxpx / denom_x;
+                cv(2,1) = cv(1,2);
+                cv(2,2) = lambdaPx*lambdaPx / denom_x;
+
+                amrex::ParticleReal denom_y = 1.0 - muypy*muypy;
+                cv(3,3) = lambdaY*lambdaY / denom_y;
+                cv(3,4) = -lambdaY*lambdaPy*muypy / denom_y;
+                cv(4,3) = cv(3,4);
+                cv(4,4) = lambdaPy*lambdaPy / denom_y;
+
+                amrex::ParticleReal denom_t = 1.0 - mutpt*mutpt;
+                cv(5,5) = lambdaT*lambdaT / denom_t;
+                cv(5,6) = -lambdaT*lambdaPt*mutpt / denom_t;
+                cv(6,5) = cv(5,6);
+                cv(6,6) = lambdaPt*lambdaPt / denom_t;
+
+            }
+        }, distr);
+
+        return cv;
+    }
+
     void
     ImpactX::add_particles (
         amrex::ParticleReal bunch_charge,
@@ -95,7 +155,7 @@ namespace impactx
             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>
@@ -436,6 +437,23 @@ namespace detail
                     read_element(sub_element_name, m_lattice, nslice_default, mapsteps_default);
                 }
             }
+        } else if (element_type == "linear_map")
+        {
+            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;
+                }
+            }
+            std::cout << "Caution: A user-provided linear map is used.  Transport may not be symplectic." << "\n";
+
+            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_REAL.H>
+#include <AMReX_SmallMatrix.H>
+
+
+namespace impactx
+{
+    /** this is a 6x6 matrix */
+    using CovarianceMatrix = amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1>;
+
+} // namespace impactx::distribution
+
+#endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H

src/particles/PushAll.H

@@ -51,6 +51,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/elements/All.H

@@ -20,20 +20,21 @@
 #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 "Marker.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"
@@ -61,6 +62,7 @@ namespace impactx
         ExactDrift,
         ExactSbend,
         Kicker,
+        LinearMap,
         Marker,
         Multipole,
         NonlinearLens,

src/particles/elements/Drift.H

@@ -16,9 +16,11 @@
 #include "mixin/thick.H"
 #include "mixin/named.H"
 #include "mixin/nofinalize.H"
+#include "mixin/lineartransport.H"
 
 #include <AMReX_Extension.H>
 #include <AMReX_REAL.H>
+#include <AMReX_SmallMatrix.H>
 
 #include <cmath>
 
@@ -160,6 +162,40 @@ namespace impactx
             refpart.s = s + slice_ds;
 
         }
+
+        /** This function returns the linear transport map.
+         *
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+
+        amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1>
+        transport_map(RefPart & AMREX_RESTRICT refpart) {
+
+            using namespace amrex::literals; // for _rt and _prt
+            elements::LinearTransport::Map6x6 R = {};
+
+            // 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;
+
+            // assign linear map matrix elements
+            R(1,1) = 1.0_prt;
+            R(1,2) = slice_ds;
+            R(2,2) = 1.0_prt;
+            R(3,3) = 1.0_prt;
+            R(3,4) = slice_ds;
+            R(4,4) = 1.0_prt;
+            R(5,5) = 1.0_prt;
+            R(5,6) = slice_ds/betgam2;
+            R(6,6) = 1.0_prt;
+            return R;
+
+        }
+
     };
 
 } // namespace impactx

src/particles/elements/LinearMap.H

@@ -0,0 +1,135 @@
+/* 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 a user-provided linear transport map R
+         *  to the 6-vector of phase space coordinates (x,px,y,py,t,pt).
+         *  Thus  x_final = R(1,1)*x + R(1,2)*px + R(1,3)*y + ...,
+         *       px_final = R(2,1)*x + R(2,2)*px + R(2,3)*y + ..., etc.
+         *
+         * @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 LinearMap functor, so that a variable of this type can be used like a
+         *  LinearMap 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);
+
+            // input and output phase space vectors
+            amrex::Array1D<amrex::ParticleReal,1,6> vectorin;
+            amrex::Array1D<amrex::ParticleReal,1,6> vectorout;
+            vectorin(1) = x;
+            vectorin(2) = px;
+            vectorin(3) = y;
+            vectorin(4) = py;
+            vectorin(5) = t;
+            vectorin(6) = pt;
+
+            for (int i=1; i<=6; ++i) {
+
+                vectorout(i) = 0.0;
+                for (int j = 1; j <= 6; ++j) {
+                    vectorout(i) += m_transport_map(i, j) * vectorin(j);
+                }
+
+            }
+
+            // assign updated values
+            x = vectorout(1);
+            px = vectorout(2);
+            y = vectorout(3);
+            py = vectorout(4);
+            t = vectorout(5);
+            pt = vectorout(6);
+
+            // undo shift due to alignment errors of the element
+            shift_out(x, y, px, py);
+        }
+
+        /** This pushes the reference particle. */
+        using Thin::operator();
+
+        LinearTransport::Map6x6 m_transport_map;  // 6x6 transport map
+
+    };
+
+} // namespace impactx
+
+#endif // IMPACTX_ELEMENT_LINEAR_MAP_H