[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_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

@@ -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/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,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,6 +16,7 @@
 #include "mixin/thick.H"
 #include "mixin/named.H"
 #include "mixin/nofinalize.H"
+#include "mixin/lineartransport.H"
 
 #include <AMReX_Extension.H>
 #include <AMReX_REAL.H>
@@ -160,6 +161,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::Array2D<amrex::ParticleReal, 1, 6, 1, 6>
+        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