From 4e2a733373002399086f5a60aa37fb4796ef5371 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 24 Sep 2024 15:33:24 -0700 Subject: [PATCH 01/10] [Draft] Transport and Covariance Matrices General structure. Co-authored-by: Chad Mitchell --- src/initialization/InitDistribution.cpp | 45 +++++++- src/initialization/InitElement.cpp | 20 ++++ src/particles/CovarianceMatrix.H | 24 ++++ src/particles/PushAll.H | 4 + src/particles/distribution/Gaussian.H | 1 - src/particles/distribution/KVdist.H | 1 - src/particles/distribution/Kurth4D.H | 1 - src/particles/distribution/Kurth6D.H | 1 - src/particles/distribution/Semigaussian.H | 1 - src/particles/distribution/Triangle.H | 2 +- src/particles/distribution/Waterbag.H | 1 - src/particles/elements/All.H | 6 +- src/particles/elements/LinearMap.H | 104 ++++++++++++++++++ .../elements/mixin/lineartransport.H | 50 +++++++++ 14 files changed, 251 insertions(+), 10 deletions(-) create mode 100644 src/particles/CovarianceMatrix.H create mode 100644 src/particles/elements/LinearMap.H create mode 100644 src/particles/elements/mixin/lineartransport.H diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 61cdd77aa..019a79729 100644 --- a/src/initialization/InitDistribution.cpp +++ b/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::value || + std::is_same::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 @@ 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 >; + using Distribution = std::remove_reference_t< std::remove_cv_t >; // TODO: switch order ov remove_ ...? initialization::InitSingleParticleData 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); diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 9c5afd6b6..7fd9671cb 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -9,6 +9,7 @@ */ #include "ImpactX.H" #include "particles/elements/All.H" +#include "particles/elements/mixin/lineartransport.H" #include #include @@ -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); } diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H new file mode 100644 index 000000000..60bd15df4 --- /dev/null +++ b/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 +#include + + +namespace impactx +{ + /** this is a 6x6 matrix */ + using CovarianceMatrix = amrex::Array2D; + +} // namespace impactx::distribution + +#endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H diff --git a/src/particles/PushAll.H b/src/particles/PushAll.H index ab0bacd65..bd3cef0ed 100644 --- a/src/particles/PushAll.H +++ b/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) diff --git a/src/particles/distribution/Gaussian.H b/src/particles/distribution/Gaussian.H index d6baf9640..56c301eb4 100644 --- a/src/particles/distribution/Gaussian.H +++ b/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 diff --git a/src/particles/distribution/KVdist.H b/src/particles/distribution/KVdist.H index 87e343399..f09193ae4 100644 --- a/src/particles/distribution/KVdist.H +++ b/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 diff --git a/src/particles/distribution/Kurth4D.H b/src/particles/distribution/Kurth4D.H index cd254d2f9..641c00ebf 100644 --- a/src/particles/distribution/Kurth4D.H +++ b/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 diff --git a/src/particles/distribution/Kurth6D.H b/src/particles/distribution/Kurth6D.H index e3d1b8c42..8dbb5d62a 100644 --- a/src/particles/distribution/Kurth6D.H +++ b/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 diff --git a/src/particles/distribution/Semigaussian.H b/src/particles/distribution/Semigaussian.H index 63d985657..0dba93ea8 100644 --- a/src/particles/distribution/Semigaussian.H +++ b/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 diff --git a/src/particles/distribution/Triangle.H b/src/particles/distribution/Triangle.H index a1f40dc14..b0c833c0c 100644 --- a/src/particles/distribution/Triangle.H +++ b/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 diff --git a/src/particles/distribution/Waterbag.H b/src/particles/distribution/Waterbag.H index 4f32fdd57..701c0b409 100644 --- a/src/particles/distribution/Waterbag.H +++ b/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 diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index 6f7f72066..03dc67f87 100644 --- a/src/particles/elements/All.H +++ b/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, diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H new file mode 100644 index 000000000..b129ef500 --- /dev/null +++ b/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 +#include + +#include + + +namespace impactx +{ + struct LinearMap + : public elements::BeamOptic, + 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 diff --git a/src/particles/elements/mixin/lineartransport.H b/src/particles/elements/mixin/lineartransport.H new file mode 100644 index 000000000..43641f21b --- /dev/null +++ b/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 + +#include +#include +#include + + +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; + Map6x6 m_transport_map; ///< linearized map + }; + +} // namespace impactx::elements + +#endif // IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H From 5d09112ee76d92f0071d8f108d89790fb3d153e3 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Wed, 25 Sep 2024 15:36:02 -0700 Subject: [PATCH 02/10] Update InitElement.cpp Update initialization of LinearMap element. --- src/initialization/InitElement.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 7fd9671cb..55cf4bc05 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -437,11 +437,8 @@ namespace detail } } 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) { @@ -452,6 +449,7 @@ namespace detail 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 { From 7d95ca57b60eae4cd339c8711f13a6290b933e71 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 25 Sep 2024 22:37:01 +0000 Subject: [PATCH 03/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/initialization/InitElement.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 55cf4bc05..36c69e633 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -438,7 +438,7 @@ namespace detail } 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) { From 65c9bd00f4b46d98539b07e6307ed11f42abebf1 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Wed, 25 Sep 2024 15:45:31 -0700 Subject: [PATCH 04/10] Update LinearMap.H Finish implementation of LinearMap element push. --- src/particles/elements/LinearMap.H | 39 +++++++++++++++++++++++++++--- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index b129ef500..98883f91f 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -35,7 +35,10 @@ namespace impactx static constexpr auto type = "LinearMap"; using PType = ImpactXParticleContainer::ParticleType; - /** A thin element that applies ... + /** 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 @@ -60,8 +63,8 @@ namespace impactx /** Push all particles */ using BeamOptic::operator(); - /** This is a ... functor, so that a variable of this type can be used like a - * function. + /** 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 @@ -89,7 +92,32 @@ namespace impactx // shift due to alignment errors of the element shift_in(x, y, px, py); - // ... TODO ... + // input and output phase space vectors + amrex::Array1D vectorin; + amrex::Array1D 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); @@ -97,6 +125,9 @@ namespace impactx /** This pushes the reference particle. */ using Thin::operator(); + + LinearTransport::Map6x6 m_transport_map; // 6x6 transport map + }; } // namespace impactx From 2ad47d2530e772e076d92a2091e1a7a970b52cc2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 25 Sep 2024 22:45:38 +0000 Subject: [PATCH 05/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/elements/LinearMap.H | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index 98883f91f..f50e765e9 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -36,8 +36,8 @@ namespace impactx 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 + ..., + * 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 @@ -94,21 +94,21 @@ namespace impactx // input and output phase space vectors amrex::Array1D vectorin; - amrex::Array1D vectorout; + amrex::Array1D vectorout; vectorin(1) = x; vectorin(2) = px; vectorin(3) = y; vectorin(4) = py; - vectorin(5) = t; + 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 From c2fe2be89b271d036d0b6100eb8636baad4c4a3d Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Wed, 25 Sep 2024 17:17:41 -0700 Subject: [PATCH 06/10] Update InitDistribution.cpp Initialize covariance matrix from distribution parameters. --- src/initialization/InitDistribution.cpp | 35 ++++++++++++++++++------- 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 019a79729..f11dc66ee 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -60,15 +60,32 @@ namespace impactx 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 + 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); From a5df0a32a439702b8f0dade266dc5193cb4c948a Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 26 Sep 2024 14:55:26 -0700 Subject: [PATCH 07/10] Add Drift return linear map. --- src/particles/elements/Drift.H | 36 ++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index c04c6d230..25d78736c 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -15,6 +15,7 @@ #include "mixin/beamoptic.H" #include "mixin/thick.H" #include "mixin/nofinalize.H" +#include "mixin/lineartransport.H" #include #include @@ -155,6 +156,41 @@ 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 + transport_map(RefPart & AMREX_RESTRICT refpart) { + + using namespace amrex::literals; // for _rt and _prt + + amrex::Array2D 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 From b2f158c9e372b676d6542221c3dee6f8c2dd9121 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Thu, 26 Sep 2024 17:43:57 -0700 Subject: [PATCH 08/10] Use LinearMap struct in Drift.H. --- src/particles/elements/Drift.H | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index 25d78736c..3e7f633f5 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -167,8 +167,7 @@ namespace impactx transport_map(RefPart & AMREX_RESTRICT refpart) { using namespace amrex::literals; // for _rt and _prt - - amrex::Array2D R; + elements::LinearTransport::Map6x6 R = {}; // length of the current slice amrex::ParticleReal const slice_ds = m_ds / nslice(); From cd4c9896d672d398ade2712648506d5a12de318f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 10 Oct 2024 17:18:32 +0000 Subject: [PATCH 09/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/elements/Drift.H | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index b6f37d007..8bf328840 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -163,20 +163,20 @@ namespace impactx } /** This function returns the linear transport map. - * + * * @returns 6x6 transport matrix */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::Array2D 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; @@ -192,7 +192,7 @@ namespace impactx R(5,6) = slice_ds/betgam2; R(6,6) = 1.0_prt; return R; - + } }; From 817cc51b88f83fd16ce23ac2f210dee6540b66fa Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 11 Oct 2024 16:16:11 -0700 Subject: [PATCH 10/10] Transport/Covariane: Use `SmallMatrix` --- src/particles/CovarianceMatrix.H | 4 ++-- src/particles/elements/Drift.H | 3 ++- src/particles/elements/mixin/lineartransport.H | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 60bd15df4..23af3e5aa 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -10,14 +10,14 @@ #ifndef IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H #define IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H -#include #include +#include namespace impactx { /** this is a 6x6 matrix */ - using CovarianceMatrix = amrex::Array2D; + using CovarianceMatrix = amrex::SmallMatrix; } // namespace impactx::distribution diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index 8bf328840..944177558 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -20,6 +20,7 @@ #include #include +#include #include @@ -168,7 +169,7 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - amrex::Array2D + amrex::SmallMatrix transport_map(RefPart & AMREX_RESTRICT refpart) { using namespace amrex::literals; // for _rt and _prt diff --git a/src/particles/elements/mixin/lineartransport.H b/src/particles/elements/mixin/lineartransport.H index 43641f21b..46416ba90 100644 --- a/src/particles/elements/mixin/lineartransport.H +++ b/src/particles/elements/mixin/lineartransport.H @@ -17,6 +17,7 @@ #include #include #include +#include namespace impactx::elements @@ -41,7 +42,7 @@ namespace impactx::elements ~LinearTransport () = default; // 6x6 linear transport map - using Map6x6 = amrex::Array2D; + using Map6x6 = amrex::SmallMatrix; Map6x6 m_transport_map; ///< linearized map };