From 4e2a733373002399086f5a60aa37fb4796ef5371 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 24 Sep 2024 15:33:24 -0700 Subject: [PATCH 01/48] [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/48] 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/48] [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/48] 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/48] [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/48] 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/48] 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/48] 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/48] [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/48] 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 }; From 1ef158eb13d0684791d37fd1935c454f92a9b05d Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 15 Oct 2024 10:59:34 -0700 Subject: [PATCH 11/48] Add Python bindings for LinearMap element. --- src/python/elements.cpp | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index b6403e71b..b92c812e4 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1480,6 +1480,39 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_TaperedPL); + py::class_ py_LinearMap(me, "LinearMap"); + py_LinearMap + .def("__repr__", + [](LinearMap const & linearmap) { + return element_name( + linearmap, + std::make_pair("R", linearmap.m_transport_map), + ); + } + ) + .def(py::init([]( + LinearTransport::Map6x6 R, + amrex::ParticleReal dx, + amrex::ParticleReal dy, + amrex::ParticleReal rotation_degree, + std::optional name + ) + { + }), + py::arg("R"), + py::arg("unit") = "dimensionless", + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("name") = py::none(), + R"(A user-provided linear map, represented as a 6x6 transport matrix.)" + ) + .def_property("R", + [](LinearMap & linearmap) { return linearmap.m_transport_map; }, + [](LinearMap & linearmap, amrex::ParticleReal xkick) { linearmap.m_transport_map = linearmap; }, + "linear map as a 6x6 transport matrix" + ) + // freestanding push function m.def("push", &Push, From ccad983dfa4b8c32a540028f92dd4f3448166c16 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 15 Oct 2024 11:14:52 -0700 Subject: [PATCH 12/48] Update elements.cpp Remove unnecessary ,. --- src/python/elements.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index b92c812e4..80fbc5cc3 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1486,7 +1486,7 @@ void init_elements(py::module& m) [](LinearMap const & linearmap) { return element_name( linearmap, - std::make_pair("R", linearmap.m_transport_map), + std::make_pair("R", linearmap.m_transport_map) ); } ) From 2b288db47da8137da027ecf564a48e01359f9d3b Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 15 Oct 2024 11:19:42 -0700 Subject: [PATCH 13/48] Update elements.cpp Correct naming in elements.cpp binding for LinearMap. --- src/python/elements.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 80fbc5cc3..9822fdcb7 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1509,7 +1509,7 @@ void init_elements(py::module& m) ) .def_property("R", [](LinearMap & linearmap) { return linearmap.m_transport_map; }, - [](LinearMap & linearmap, amrex::ParticleReal xkick) { linearmap.m_transport_map = linearmap; }, + [](LinearMap & linearmap, LinearTransport::Map6x6 R) { linearmap.m_transport_map = linearmap; }, "linear map as a 6x6 transport matrix" ) From 06718d84fcac188bc56b6ffebd260f91caf403bd Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Wed, 23 Oct 2024 15:15:28 -0700 Subject: [PATCH 14/48] Update elements.cpp Correct argument declaration for Python bindings of LinearMap in elements.cpp. --- src/python/elements.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 9822fdcb7..747cde617 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1491,7 +1491,7 @@ void init_elements(py::module& m) } ) .def(py::init([]( - LinearTransport::Map6x6 R, + elements::LinearTransport::Map6x6 R, amrex::ParticleReal dx, amrex::ParticleReal dy, amrex::ParticleReal rotation_degree, From 2e57e9159ec5ba4332b4349375f020e92f49f6de Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Sat, 2 Nov 2024 10:14:40 -0700 Subject: [PATCH 15/48] Add LinearMap to 'Named' elements. --- src/particles/elements/LinearMap.H | 10 +++++++--- src/python/elements.cpp | 25 ++++++++++++------------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index f50e765e9..87bc360e0 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -15,6 +15,7 @@ #include "mixin/beamoptic.H" #include "mixin/lineartransport.H" #include "mixin/thin.H" +#include "mixin/named.H" #include "mixin/nofinalize.H" #include @@ -26,7 +27,8 @@ namespace impactx { struct LinearMap - : public elements::BeamOptic, + : public elements::Named, + public elements::BeamOptic, public elements::Thin, public elements::Alignment, public elements::LinearTransport, @@ -49,9 +51,11 @@ namespace impactx LinearTransport::Map6x6 const & R, amrex::ParticleReal dx = 0, amrex::ParticleReal dy = 0, - amrex::ParticleReal rotation_degree = 0 + amrex::ParticleReal rotation_degree = 0, + std::optional name = std::nullopt ) - : Alignment(dx, dy, rotation_degree) + : Named(name), + Alignment(dx, dy, rotation_degree) { for (int i=1; i<=6; ++i) { for (int j = 1; j <= 6; ++j) { diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 747cde617..e1f7e1f19 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -1480,7 +1481,7 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_TaperedPL); - py::class_ py_LinearMap(me, "LinearMap"); + py::class_ py_LinearMap(me, "LinearMap"); py_LinearMap .def("__repr__", [](LinearMap const & linearmap) { @@ -1490,17 +1491,14 @@ void init_elements(py::module& m) ); } ) - .def(py::init([]( - elements::LinearTransport::Map6x6 R, - amrex::ParticleReal dx, - amrex::ParticleReal dy, - amrex::ParticleReal rotation_degree, - std::optional name - ) - { - }), + .def(py::init< + elements::LinearTransport::Map6x6, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + std::optional + >(), py::arg("R"), - py::arg("unit") = "dimensionless", py::arg("dx") = 0, py::arg("dy") = 0, py::arg("rotation") = 0, @@ -1509,10 +1507,11 @@ void init_elements(py::module& m) ) .def_property("R", [](LinearMap & linearmap) { return linearmap.m_transport_map; }, - [](LinearMap & linearmap, LinearTransport::Map6x6 R) { linearmap.m_transport_map = linearmap; }, + [](LinearMap & linearmap, elements::LinearTransport::Map6x6 R) { linearmap.m_transport_map = R; }, "linear map as a 6x6 transport matrix" ) - + ; + register_beamoptics_push(py_LinearMap); // freestanding push function m.def("push", &Push, From 911fbeca093754482bcfe38dc146fed562c6059a Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Sat, 2 Nov 2024 10:18:11 -0700 Subject: [PATCH 16/48] Update LinearMap.H Document "Name" parameter in LinearMap. --- src/particles/elements/LinearMap.H | 1 + 1 file changed, 1 insertion(+) diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index 87bc360e0..0bbce8a14 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -46,6 +46,7 @@ namespace impactx * @param dx horizontal translation error in m * @param dy vertical translation error in m * @param rotation_degree rotation error in the transverse plane [degrees] + * @param name a user defined and not necessarily unique name of the element */ LinearMap ( LinearTransport::Map6x6 const & R, From ac6556c0381a0b3d52eaafd2c0676ebbfb2d9899 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Sat, 2 Nov 2024 11:03:05 -0700 Subject: [PATCH 17/48] Modify LinearMap Python bindings. --- src/python/elements.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index e1f7e1f19..f25c589a6 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1487,7 +1487,8 @@ void init_elements(py::module& m) [](LinearMap const & linearmap) { return element_name( linearmap, - std::make_pair("R", linearmap.m_transport_map) + std::make_pair("R11", linearmap.m_transport_map(1,1)), + std::make_pair("R12", linearmap.m_transport_map(1,2)) ); } ) @@ -1503,7 +1504,7 @@ void init_elements(py::module& m) py::arg("dy") = 0, py::arg("rotation") = 0, py::arg("name") = py::none(), - R"(A user-provided linear map, represented as a 6x6 transport matrix.)" + "(A user-provided linear map, represented as a 6x6 transport matrix.)" ) .def_property("R", [](LinearMap & linearmap) { return linearmap.m_transport_map; }, From 0444b5d0f5ec0ed815ca94203ed88cfaa89f9402 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 5 Nov 2024 11:42:54 -0800 Subject: [PATCH 18/48] Update src/initialization/InitDistribution.cpp --- src/initialization/InitDistribution.cpp | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index f11dc66ee..11be408cc 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -42,11 +42,7 @@ namespace impactx { // 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; - } - } +cv.setVal(0.0); // initialize from 2nd order beam moments std::visit([&](auto&& distribution) { From fa6792183e4d728072d7ead1cdfdb2b191b203e3 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 5 Nov 2024 11:44:00 -0800 Subject: [PATCH 19/48] Update src/initialization/InitDistribution.cpp --- src/initialization/InitDistribution.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 11be408cc..d1a3aae34 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -41,8 +41,7 @@ namespace impactx ) { // zero out the 6x6 matrix - CovarianceMatrix cv; -cv.setVal(0.0); + CovarianceMatrix cv(0); // initialize from 2nd order beam moments std::visit([&](auto&& distribution) { From d460272874202fd3029a18709adc3073c924ae3e Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Wed, 20 Nov 2024 15:27:53 -0800 Subject: [PATCH 20/48] Update InitDistribution.cpp Modify zero-initialization of CovarianceMatrix cv. --- src/initialization/InitDistribution.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index d1a3aae34..1b848a6e1 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -41,7 +41,7 @@ namespace impactx ) { // zero out the 6x6 matrix - CovarianceMatrix cv(0); + CovarianceMatrix cv{}; // initialize from 2nd order beam moments std::visit([&](auto&& distribution) { From 9fb2a2ae37990c29dde9c96edf70f81e1ab3d35c Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 5 Jan 2025 21:30:44 -0800 Subject: [PATCH 21/48] Leftover Improvements from first PR --- src/initialization/InitElement.cpp | 17 ++++++++++------ src/particles/elements/Drift.H | 25 ++++++++--------------- src/particles/elements/LinearMap.H | 32 ++++++++---------------------- 3 files changed, 27 insertions(+), 47 deletions(-) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index efbcd0390..def524f4f 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -11,6 +11,8 @@ #include "particles/elements/All.H" #include "particles/elements/mixin/lineartransport.H" +#include + #include #include #include @@ -441,17 +443,20 @@ namespace detail { auto a = detail::query_alignment(pp_element); - elements::LinearTransport::Map6x6 transport_map; + elements::LinearTransport::Map6x6 transport_map = elements::LinearTransport::Map6x6::Identity(); + + // safe to ParmParse inputs for reproducibility 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; + pp_element.queryAdd(name.c_str(), transport_map(i, j)); } } - std::cout << "Caution: A user-provided linear map is used. Transport may not be symplectic." << "\n"; + ablastr::warn_manager::WMRecordWarning( + "ImpactX::read_element", + "Caution, a user-provided linear map is used. Transport may not be symplectic.", + ablastr::warn_manager::WarnPriority::low + ); m_lattice.emplace_back(LinearMap(transport_map, a["dx"], a["dy"], a["rotation_degree"]) ); } else { diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index 944177558..d29e95709 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -131,8 +131,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 @@ -160,7 +160,6 @@ namespace impactx // advance integrated path length refpart.s = s + slice_ds; - } /** This function returns the linear transport map. @@ -168,12 +167,10 @@ namespace impactx * @returns 6x6 transport matrix */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - - amrex::SmallMatrix - transport_map(RefPart & AMREX_RESTRICT refpart) { - + elements::LinearTransport::Map6x6 + transport_map (RefPart & AMREX_RESTRICT refpart) const + { 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(); @@ -183,19 +180,13 @@ namespace impactx amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; // assign linear map matrix elements - R(1,1) = 1.0_prt; + elements::LinearTransport::Map6x6 R = elements::LinearTransport::Map6x6::Identity(); 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; + R(5,6) = slice_ds / betgam2; + return R; } - }; } // namespace impactx diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index 0bbce8a14..99a5e3a41 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -55,14 +55,10 @@ namespace impactx amrex::ParticleReal rotation_degree = 0, std::optional name = std::nullopt ) - : Named(name), - Alignment(dx, dy, rotation_degree) + : Named(std::move(name)), + 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); - } - } + m_transport_map = R; } /** Push all particles */ @@ -98,23 +94,11 @@ namespace impactx shift_in(x, y, px, py); // 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); - } - - } + amrex::SmallVector vectorin{ + x, px, y, py, t, pt + }; + + amrex::SmallVector const vectorout = m_transport_map * vectorin; // assign updated values x = vectorout(1); From f1a13af911ac15b73bab5618e6932ff914f9bda4 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 5 Jan 2025 23:58:08 -0800 Subject: [PATCH 22/48] Finalize Python --- src/python/elements.cpp | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index f25c589a6..8e54176d0 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -211,6 +211,22 @@ void init_elements(py::module& m) ) ; + py::class_(me, "LinearTransport") + .def(py::init<>(), + "Mixin class for linear transport approximation via matrices." + ) + // type of map + .def_property_readonly_static("Map6x6", + [](py::object /* lt */){ return py::type::of(); }, + "1-indexed, Fortran-ordered, 6x6 linear transport map type" + ) + // values of the map + .def_property_readonly("transport_map", + [](elements::LinearTransport const & lt) { return lt.m_transport_map; }, + "1-indexed, Fortran-ordered, 6x6 linear transport map values" + ) + ; + // diagnostics py::class_ py_BeamMonitor(me, "BeamMonitor"); From 0807b4156d19fd36483103ccd207812d53b77a39 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 5 Jan 2025 23:58:36 -0800 Subject: [PATCH 23/48] More Detailed Warning Include element name in output to user. --- src/initialization/InitElement.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index def524f4f..1cabe432e 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -454,7 +454,8 @@ namespace detail } ablastr::warn_manager::WMRecordWarning( "ImpactX::read_element", - "Caution, a user-provided linear map is used. Transport may not be symplectic.", + "Caution, a user-provided linear map is used in '" + + element_name + "'. Transport may not be symplectic.", ablastr::warn_manager::WarnPriority::low ); From f7cb2e2853a0242467b37e6a87db0ba1433602fd Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 5 Jan 2025 23:59:41 -0800 Subject: [PATCH 24/48] Start Examples --- examples/CMakeLists.txt | 17 ++++- examples/fodo_userdef/README.rst | 85 +++++++++++++++++++++ examples/fodo_userdef/analysis_fodo.py | 1 + examples/fodo_userdef/input_fodo_userdef.in | 59 ++++++++++++++ examples/fodo_userdef/plot_fodo.py | 1 + examples/fodo_userdef/run_fodo_userdef.py | 85 +++++++++++++++++++++ 6 files changed, 247 insertions(+), 1 deletion(-) create mode 100644 examples/fodo_userdef/README.rst create mode 120000 examples/fodo_userdef/analysis_fodo.py create mode 100644 examples/fodo_userdef/input_fodo_userdef.in create mode 120000 examples/fodo_userdef/plot_fodo.py create mode 100755 examples/fodo_userdef/run_fodo_userdef.py diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 19c4d33ca..f2cce944e 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -82,7 +82,7 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) amrex.throw_exception = 1 amrex.signal_handling = 0 impactx.always_warn_immediately=1 - impactx.abort_on_warning_threshold=low + impactx.abort_on_warning_threshold=medium WORKING_DIRECTORY ${THIS_WORKING_DIR} ) endif() @@ -210,6 +210,21 @@ add_impactx_test(FODO.MADX.py examples/fodo/plot_fodo.py ) +# Python: FODO Cell w/ custom linear element ################################# +# +add_impactx_test(FODO.userdef + examples/fodo_userdef/input_fodo_userdef.in + ON # ImpactX MPI-parallel + examples/fodo_userdef/analysis_fodo.py + examples/fodo_userdef/plot_fodo.py +) +add_impactx_test(FODO.userdef.py + examples/fodo_userdef/run_fodo_userdef.py + OFF # ImpactX MPI-parallel + examples/fodo_userdef/analysis_fodo.py + examples/fodo_userdef/plot_fodo.py +) + # Python: MPI-parallel FODO Cell ############################################## # add_impactx_test(FODO.py.MPI diff --git a/examples/fodo_userdef/README.rst b/examples/fodo_userdef/README.rst new file mode 100644 index 000000000..657c94071 --- /dev/null +++ b/examples/fodo_userdef/README.rst @@ -0,0 +1,85 @@ +.. _examples-fodo-userdef: + +User-Defined Linear Element +=========================== + +This implements the same FODO cell as the :ref:`stable FODO cell example `. +However, in the example here we define *additional user-defined, custom linear elements* by providing a custom matrix. + +.. note:: + + Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. + For an even more general, user-defined element, see :ref:`the FODO Cell example that uses a Programmable Element `. + For more details, see :ref:`this section `. + +The matched Twiss parameters at entry are: + +* :math:`\beta_\mathrm{x} = 2.82161941` m +* :math:`\alpha_\mathrm{x} = -1.59050035` +* :math:`\beta_\mathrm{y} = 2.82161941` m +* :math:`\alpha_\mathrm{y} = 1.59050035` + +We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm. + +The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling. + +In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_fodo_userdef.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo_userdef.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_fodo_userdef.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_userdef/run_fodo_userdef.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo_userdef.in + :language: ini + :caption: You can copy this file from ``examples/fodo_userdef/input_fodo_userdef.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo.py`` + + .. literalinclude:: analysis_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_userdef/analysis_fodo.py``. + + +Visualize +--------- + +You can run the following script to visualize the beam evolution over time: + +.. dropdown:: Script ``plot_fodo.py`` + + .. literalinclude:: plot_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_userdef/plot_fodo.py``. + +.. figure:: https://user-images.githubusercontent.com/1353258/180287840-8561f6fd-278f-4856-abd8-04fbdb78c8ff.png + :alt: focusing, defocusing and preserved emittance in our FODO cell benchmark. + + FODO transversal beam width and emittance evolution + +.. figure:: https://user-images.githubusercontent.com/1353258/180287845-eb0210a7-2500-4aa9-844c-67fb094329d3.png + :alt: focusing, defocusing and phase space rotation in our FODO cell benchmark. + + FODO transversal beam width and phase space evolution diff --git a/examples/fodo_userdef/analysis_fodo.py b/examples/fodo_userdef/analysis_fodo.py new file mode 120000 index 000000000..dc8eb9737 --- /dev/null +++ b/examples/fodo_userdef/analysis_fodo.py @@ -0,0 +1 @@ +../fodo/analysis_fodo.py \ No newline at end of file diff --git a/examples/fodo_userdef/input_fodo_userdef.in b/examples/fodo_userdef/input_fodo_userdef.in new file mode 100644 index 000000000..0686b5fbf --- /dev/null +++ b/examples/fodo_userdef/input_fodo_userdef.in @@ -0,0 +1,59 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 2.0e3 +beam.charge = 1.0e-9 +beam.particle = electron +beam.distribution = waterbag_from_twiss +beam.alphaX = -1.5905003499999992 +beam.alphaY = 1.5905003499999992 +beam.alphaT = 0.0 +beam.betaX = 2.8216194100262637 +beam.betaY = 2.8216194100262637 +beam.betaT = 0.5 +beam.emittX = 2e-09 +beam.emittY = 2e-09 +beam.emittT = 2e-06 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift1 monitor +lattice.nslice = 25 + +monitor.type = beam_monitor +monitor.backend = h5 + +drift1.type = linear_map +drift1.R12 = 0.25 # ds +drift1.R34 = 0.25 # ds +drift1.R56 = 0.25 / 16.6464 # ds / (beta*gamma^2) + +quad1.type = quad +quad1.ds = 1.0 +quad1.k = 1.0 + +drift2.type = linear_map +drift2.R12 = 0.5 # ds +drift2.R34 = 0.5 # ds +drift2.R56 = 0.5 / 16.6464 # ds / (beta*gamma^2) + +quad2.type = quad +quad2.ds = 1.0 +quad2.k = -1.0 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/fodo_userdef/plot_fodo.py b/examples/fodo_userdef/plot_fodo.py new file mode 120000 index 000000000..ce8494b77 --- /dev/null +++ b/examples/fodo_userdef/plot_fodo.py @@ -0,0 +1 @@ +../fodo/plot_fodo.py \ No newline at end of file diff --git a/examples/fodo_userdef/run_fodo_userdef.py b/examples/fodo_userdef/run_fodo_userdef.py new file mode 100755 index 000000000..6f425db05 --- /dev/null +++ b/examples/fodo_userdef/run_fodo_userdef.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2024 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Marco Garten +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements, twiss + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 2.0e3 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 10000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Waterbag( + **twiss( + beta_x=2.8216194100262637, + beta_y=2.8216194100262637, + beta_t=0.5, + emitt_x=2e-09, + emitt_y=2e-09, + emitt_t=2e-06, + alpha_x=-1.5905003499999992, + alpha_y=1.5905003499999992, + alpha_t=0.0, + ) +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# add a user-defined, linear element for the drifts +Iden = elements.LinearMap.Map6x6.identity() +R1, R2 = Iden, Iden +R1[1, 2] = 0.25 # ds +R1[3, 4] = 0.25 # ds +R1[5, 6] = 0.25 / 16.6464 # ds / (beta*gamma^2) +drift1 = elements.LinearMap(name="drift1", R=R1) +R2[1, 2] = 0.5 # ds +R2[3, 4] = 0.5 # ds +R2[5, 6] = 0.5 / 16.6464 # ds / (beta*gamma^2) +drift2 = elements.LinearMap(name="drift2", R=R2) + +# design the accelerator lattice) +ns = 25 # number of slices per ds in the element +fodo = [ + monitor, + drift1, + monitor, + elements.Quad(name="quad1", ds=1.0, k=1.0, nslice=ns), + monitor, + drift2, + monitor, + elements.Quad(name="quad2", ds=1.0, k=-1.0, nslice=ns), + monitor, + drift1, + monitor, +] +# assign a fodo segment +sim.lattice.extend(fodo) + +# run simulation +sim.evolve() + +# clean shutdown +sim.finalize() From e1706ebe957978feca9dc4d43c914719da72f359 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 5 Jan 2025 23:59:56 -0800 Subject: [PATCH 25/48] Start Documentation --- docs/source/usage/parameters.rst | 15 +++++++++++++++ docs/source/usage/python.rst | 12 ++++++++++++ docs/source/usage/workflows/add_element.rst | 10 ++++++++-- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 1b0f2644b..59ac81cbf 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -349,6 +349,21 @@ Lattice Elements * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane * ``.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``) + * ``linear_map`` for a custom, linear transport matrix. + + The matrix elements are 1-indexed. + The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that are differing from that need to be specified. + + TODO: describe units to put in the R entries. + + This requires these additional parameters: + + * ``.R(i,j)`` (``float``, ...) matrix entries + a 1-indexed, 6x6, linear transport map to multiply with the the beam vector :math:`x,px,y,py,t,pt`. + * ``.dx`` (``float``, in meters) horizontal translation error + * ``.dy`` (``float``, in meters) vertical translation error + * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane + * ``multipole`` for a thin multipole element. This requires these additional parameters: diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 73e9c0329..85607b6d1 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -607,6 +607,18 @@ This module provides elements for the accelerator lattice. :param unit: specification of units (``"dimensionless"`` in units of the magnetic rigidity of the reference particle or ``"T-m"``) :param name: an optional name for the element +.. py::class:: impactx.elements.LinearMap(R, dx=0, dy=0, rotation=0, name=None) + + A custom, linear transport matrix. + + TODO: describe units to put in the R entries. + + :param R: a linear transport map to multiply with the the beam vector :math:`x,px,y,py,t,pt`. + :param dx: horizontal translation error in m + :param dy: vertical translation error in m + :param rotation: rotation error in the transverse plane [degrees] + :param name: an optional name for the element + .. py:class:: impactx.elements.Multipole(multipole, K_normal, K_skew, dx=0, dy=0, rotation=0, name=None) A general thin multipole element. diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index db80e21eb..106284ff1 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -33,10 +33,16 @@ Detailed particle computing interfaces are presented in the `pyAMReX examples ` for Python and inputs file syntax to specify a custom linear element. + +The entries of the transport matrix are ordered for multiplication with the beam vector :math:`x,px,y,py,t,pt`. +The units of the transport matrix :math:`R` elements are... TODO + .. note:: - We plan to add a simple, linear map element that can be configured in user input. - Follow `issue #538 `__ for progress. + Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. + Thus, where symplecticity is needed, the programmable element and the C++ element provide a more general approach. C++ Element From 26f45dba1360e1c0290385187a95272972ba0d0e Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 6 Jan 2025 00:23:28 -0800 Subject: [PATCH 26/48] Fix Bugs (incl. AMReX) Use parser for inputs. Needs upstream AMReX fix. --- src/initialization/InitElement.cpp | 2 +- src/particles/elements/LinearMap.H | 1 - src/particles/elements/mixin/lineartransport.H | 3 ++- src/python/elements.cpp | 8 ++++---- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 1cabe432e..c561957d6 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -449,7 +449,7 @@ namespace detail for (int i=1; i<=6; ++i) { for (int j=1; j<=6; ++j) { std::string name = "R" + std::to_string(i) + std::to_string(j); - pp_element.queryAdd(name.c_str(), transport_map(i, j)); + pp_element.queryAddWithParser(name.c_str(), transport_map(i, j)); } } ablastr::warn_manager::WMRecordWarning( diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index 99a5e3a41..d3bda4379 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -116,7 +116,6 @@ namespace impactx using Thin::operator(); LinearTransport::Map6x6 m_transport_map; // 6x6 transport map - }; } // namespace impactx diff --git a/src/particles/elements/mixin/lineartransport.H b/src/particles/elements/mixin/lineartransport.H index 46416ba90..9fce44ff6 100644 --- a/src/particles/elements/mixin/lineartransport.H +++ b/src/particles/elements/mixin/lineartransport.H @@ -43,7 +43,8 @@ namespace impactx::elements // 6x6 linear transport map using Map6x6 = amrex::SmallMatrix; - Map6x6 m_transport_map; ///< linearized map + // note: for most elements, R is returned by a member function. Some store it also internally as a member. + // Map6x6 m_transport_map; ///< linearized map }; } // namespace impactx::elements diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 8e54176d0..f3930fd72 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -221,10 +221,10 @@ void init_elements(py::module& m) "1-indexed, Fortran-ordered, 6x6 linear transport map type" ) // values of the map - .def_property_readonly("transport_map", - [](elements::LinearTransport const & lt) { return lt.m_transport_map; }, - "1-indexed, Fortran-ordered, 6x6 linear transport map values" - ) + //.def_property_readonly("R", + // [](elements::LinearTransport const & lt) { return lt.m_transport_map; }, + // "1-indexed, Fortran-ordered, 6x6 linear transport map values" + //) ; // diagnostics From a0e147f10c7640073d543d023ed43c5c5396dcca Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 6 Jan 2025 00:28:43 -0800 Subject: [PATCH 27/48] Docs Updates --- docs/source/usage/examples.rst | 1 + docs/source/usage/workflows/add_element.rst | 34 ++++++++++++--------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 0658cb278..242de3195 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -31,6 +31,7 @@ Single Particle Dynamics examples/aperture/README.rst examples/iota_lens/README.rst examples/achromatic_spectrometer/README.rst + examples/fodo_userdef/README.rst examples/fodo_programmable/README.rst examples/dogleg/README.rst examples/coupled_optics/README.rst diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index 106284ff1..3610d8239 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -10,6 +10,25 @@ The workflows described here apply both for thin kicks or thick elements. Thick elements can also use soft-edged fringe fields (see `existing soft-edged elements for implementation details `__). +.. _usage-workflows-add-element-linmap: + +Linear Map +---------- + +A custom linear element can be provided by specifying the 6x6 linear transport matrix :math:`R` as an input. +See the :ref:` example ` for Python and inputs file syntax to specify a custom linear element. + +The entries of the transport matrix are ordered for multiplication with the beam vector :math:`x,px,y,py,t,pt`. +The units of the transport matrix :math:`R` elements are... TODO + +.. note:: + + Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. + Thus, where symplecticity is needed, the :ref:`Python Programmable Element ` and the :ref:`C++ Element ` provide a more general approach. + + +.. _usage-workflows-add-element-python: + Python Programmable Element --------------------------- @@ -30,20 +49,7 @@ Detailed examples that show usage of the programmable element are: Detailed particle computing interfaces are presented in the `pyAMReX examples `__. -Linear Map ----------- - -A custom linear element can be provided by specifying the 6x6 linear transport matrix :math:`R` as an input. -See the :ref:` example ` for Python and inputs file syntax to specify a custom linear element. - -The entries of the transport matrix are ordered for multiplication with the beam vector :math:`x,px,y,py,t,pt`. -The units of the transport matrix :math:`R` elements are... TODO - -.. note:: - - Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. - Thus, where symplecticity is needed, the programmable element and the C++ element provide a more general approach. - +.. _usage-workflows-add-element-cxx: C++ Element ----------- From 317e6bfa4260dd83f682da78d961a962360ccd80 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 9 Jan 2025 19:37:21 -0800 Subject: [PATCH 28/48] `LinearMap`: `ds` bookkeeping And reference particle pushing (drifting), if thick. --- examples/fodo_userdef/run_fodo_userdef.py | 20 ++++--- src/particles/elements/LinearMap.H | 67 +++++++++++++++++++++-- src/python/elements.cpp | 15 ++++- 3 files changed, 88 insertions(+), 14 deletions(-) diff --git a/examples/fodo_userdef/run_fodo_userdef.py b/examples/fodo_userdef/run_fodo_userdef.py index 6f425db05..dd15459ee 100755 --- a/examples/fodo_userdef/run_fodo_userdef.py +++ b/examples/fodo_userdef/run_fodo_userdef.py @@ -51,14 +51,18 @@ # add a user-defined, linear element for the drifts Iden = elements.LinearMap.Map6x6.identity() R1, R2 = Iden, Iden -R1[1, 2] = 0.25 # ds -R1[3, 4] = 0.25 # ds -R1[5, 6] = 0.25 / 16.6464 # ds / (beta*gamma^2) -drift1 = elements.LinearMap(name="drift1", R=R1) -R2[1, 2] = 0.5 # ds -R2[3, 4] = 0.5 # ds -R2[5, 6] = 0.5 / 16.6464 # ds / (beta*gamma^2) -drift2 = elements.LinearMap(name="drift2", R=R2) + +ds1 = 0.25 +R1[1, 2] = ds1 +R1[3, 4] = ds1 +R1[5, 6] = ds1 / 16.6464 # ds / (beta*gamma^2) +drift1 = elements.LinearMap(name="drift1", R=R1, ds=ds1) + +ds2 = 0.5 +R2[1, 2] = ds2 +R2[3, 4] = ds2 +R2[5, 6] = ds2 / 16.6464 # ds / (beta*gamma^2) +drift2 = elements.LinearMap(name="drift2", R=R2, ds=ds2) # design the accelerator lattice) ns = 25 # number of slices per ds in the element diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H index d3bda4379..472ae12be 100644 --- a/src/particles/elements/LinearMap.H +++ b/src/particles/elements/LinearMap.H @@ -14,7 +14,6 @@ #include "mixin/alignment.H" #include "mixin/beamoptic.H" #include "mixin/lineartransport.H" -#include "mixin/thin.H" #include "mixin/named.H" #include "mixin/nofinalize.H" @@ -29,7 +28,6 @@ namespace impactx struct LinearMap : public elements::Named, public elements::BeamOptic, - public elements::Thin, public elements::Alignment, public elements::LinearTransport, public elements::NoFinalize @@ -43,6 +41,7 @@ namespace impactx * px_final = R(2,1)*x + R(2,2)*px + R(2,3)*y + ..., etc. * * @param R user-provided transport map + * @param ds Segment length in m * @param dx horizontal translation error in m * @param dy vertical translation error in m * @param rotation_degree rotation error in the transverse plane [degrees] @@ -50,6 +49,7 @@ namespace impactx */ LinearMap ( LinearTransport::Map6x6 const & R, + amrex::ParticleReal ds = 0, amrex::ParticleReal dx = 0, amrex::ParticleReal dy = 0, amrex::ParticleReal rotation_degree = 0, @@ -59,6 +59,7 @@ namespace impactx Alignment(dx, dy, rotation_degree) { m_transport_map = R; + m_ds = ds; } /** Push all particles */ @@ -112,10 +113,68 @@ namespace impactx shift_out(x, y, px, py); } - /** This pushes the reference particle. */ - using Thin::operator(); + /** This pushes the reference particle. + * + * @param[in,out] refpart reference particle + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() (RefPart & AMREX_RESTRICT refpart) const + { + if (m_ds > 0) // Drift + { + using namespace amrex::literals; // for _rt and _prt + + // assign input reference particle values + amrex::ParticleReal const x = refpart.x; + amrex::ParticleReal const px = refpart.px; + amrex::ParticleReal const y = refpart.y; + amrex::ParticleReal const py = refpart.py; + amrex::ParticleReal const z = refpart.z; + amrex::ParticleReal const pz = refpart.pz; + amrex::ParticleReal const t = refpart.t; + amrex::ParticleReal const pt = refpart.pt; + amrex::ParticleReal const s = refpart.s; + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // assign intermediate parameter + amrex::ParticleReal const step = slice_ds /std::sqrt(std::pow(pt,2)-1.0_prt); + + // advance position and momentum (drift) + refpart.x = x + step*px; + refpart.y = y + step*py; + refpart.z = z + step*pz; + refpart.t = t - step*pt; + + // advance integrated path length + refpart.s = s + slice_ds; + } + // else nothing to do for a zero-length element + } + + /** Number of slices used for the application of space charge + * + * @return one, because we do not support slicing of this element + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + int nslice () const + { + return 1; + } + + /** Return the segment length + * + * @return by default zero, but users can set a corresponding ds for bookkeeping + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + amrex::ParticleReal ds () const + { + return m_ds; + } LinearTransport::Map6x6 m_transport_map; // 6x6 transport map + amrex::ParticleReal m_ds; // finite ds allowed for bookkeeping, but we do not allow slicing }; } // namespace impactx diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 77564be51..256c9cf5f 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -216,7 +216,7 @@ void init_elements(py::module& m) ) ; - py::class_(me, "LinearTransport") + py::class_(mx, "LinearTransport") .def(py::init<>(), "Mixin class for linear transport approximation via matrices." ) @@ -1573,7 +1573,7 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_TaperedPL); - py::class_ py_LinearMap(me, "LinearMap"); + py::class_ py_LinearMap(me, "LinearMap"); py_LinearMap .def("__repr__", [](LinearMap const & linearmap) { @@ -1589,9 +1589,11 @@ void init_elements(py::module& m) amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal, + amrex::ParticleReal, std::optional >(), py::arg("R"), + py::arg("ds") = 0, py::arg("dx") = 0, py::arg("dy") = 0, py::arg("rotation") = 0, @@ -1603,6 +1605,15 @@ void init_elements(py::module& m) [](LinearMap & linearmap, elements::LinearTransport::Map6x6 R) { linearmap.m_transport_map = R; }, "linear map as a 6x6 transport matrix" ) + .def_property("ds", + [](LinearMap & linearmap) { return linearmap.m_ds; }, + [](LinearMap & linearmap, amrex::ParticleReal ds) { linearmap.m_ds = ds; }, + "segment length in m" + ) + .def_property_readonly("nslice", + [](LinearMap & linearmap) { return linearmap.nslice(); }, + "one, because we do not support slicing of this element" + ) ; register_beamoptics_push(py_LinearMap); From 2fee7e2ca906dd9ee99bf1e6ea10e619feeeffcf Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Sat, 11 Jan 2025 17:40:36 -0800 Subject: [PATCH 29/48] Add thin linear map example. --- examples/CMakeLists.txt | 16 +++ examples/linear_map/README.rst | 58 +++++++++++ examples/linear_map/analysis_map.py | 152 ++++++++++++++++++++++++++++ examples/linear_map/input_map.in | 57 +++++++++++ examples/linear_map/run_map.py | 103 +++++++++++++++++++ 5 files changed, 386 insertions(+) create mode 100644 examples/linear_map/README.rst create mode 100755 examples/linear_map/analysis_map.py create mode 100644 examples/linear_map/input_map.in create mode 100755 examples/linear_map/run_map.py diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d4de8d072..e9f6b29d1 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1071,3 +1071,19 @@ add_impactx_test(linac-segment.py examples/linac_segment/analysis_linac_segment.py OFF # no plot script yet ) + +# Iteration of a linear one-turn map ######################################### +# +# w/o space charge +add_impactx_test(linear-map + examples/linear_map/input_map.in + ON # ImpactX MPI-parallel + examples/linear_map/analysis_map.py + OFF # no plot script yet +) +add_impactx_test(linear-map.py + examples/linear_map/run_map.py + ON # ImpactX MPI-parallel + examples/linear_map/analysis_map.py + OFF # no plot script yet +) diff --git a/examples/linear_map/README.rst b/examples/linear_map/README.rst new file mode 100644 index 000000000..fad78a4ac --- /dev/null +++ b/examples/linear_map/README.rst @@ -0,0 +1,58 @@ +.. _examples-linear-map: + +Iteration of a user-defined linear map +======================================= + +This example illustrates the application of a user-defined linear map via a matrix. + +Here, the linear map represents an abstract symplectic transformation of the beam in 6D phase space. +If desired, the user may interpret the matrix as the one-turn map of a storage ring or circular collider. + +The (fractional) tunes (Qx, Qy, Qt) of the map are given by (0.139, 0.219, 0.0250). +We use a 45.6 GeV electron beam that is invariant under the action of the linear map (matched). +The horizontal and vertical unnormalized emittances are 0.27 nm and 1.0 pm, respectively. + +These parameters are based on the single-beam parameters of FCC-ee (Z-mode) appearing here: +https://twiki.cern.ch/twiki/bin/view/FCC/FCCeeParameters_CDRBaseline-1_0 + +The second moments of the phase space variables should be unchanged under application of the map. + +In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + +In addition, the tunes associated with a single particle orbit are extracted, and must agree with the values given above. + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_multipole.py`` or +* ImpactX **executable** using an input file: ``impactx input_multipole.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_multipole.py + :language: python3 + :caption: You can copy this file from ``examples/multipole/run_multipole.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_multipole.in + :language: ini + :caption: You can copy this file from ``examples/multipole/input_multipole.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_multipole.py`` + + .. literalinclude:: analysis_multipole.py + :language: python3 + :caption: You can copy this file from ``examples/multipole/analysis_multipole.py``. diff --git a/examples/linear_map/analysis_map.py b/examples/linear_map/analysis_map.py new file mode 100755 index 000000000..86cb5ec0c --- /dev/null +++ b/examples/linear_map/analysis_map.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import numpy as np +import openpmd_api as io +from scipy.stats import moment + +def get_moments(beam): + """Calculate standard deviations of beam position & momenta + and emittance values + + Returns + ------- + sigx, sigy, sigt, emittance_x, emittance_y, emittance_t + """ + sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev. + sigpx = moment(beam["momentum_x"], moment=2) ** 0.5 + sigy = moment(beam["position_y"], moment=2) ** 0.5 + sigpy = moment(beam["momentum_y"], moment=2) ** 0.5 + sigt = moment(beam["position_t"], moment=2) ** 0.5 + sigpt = moment(beam["momentum_t"], moment=2) ** 0.5 + + epstrms = beam.cov(ddof=0) + emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5 + emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5 + emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5 + + return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t) + +# initial/final beam +series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) +last_step = list(series.iterations)[-1] +initial = series.iterations[1].particles["beam"].to_df() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 6.363961030678928e-6, + 28.284271247461902e-9, + 0.0035, + 0.27e-9, + 1.0e-12, + 1.33e-6, + ], + rtol=rtol, + atol=atol, +) + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 6.363961030678928e-6, + 28.284271247461902e-9, + 0.0035, + 0.27e-9, + 1.0e-12, + 1.33e-6, + ], + rtol=rtol, + atol=atol, +) + +# Specify time series for particle j +j = 5 +print(f"output for particle index = {j}") + +# Create array of TBT data values +x = [] +px = [] +y = [] +py = [] +t = [] +pt = [] +n = 0 +for k_i, i in series.iterations.items(): + beam = i.particles["beam"] + turn = beam.to_df() + x.append(turn["position_x"][j]) + px.append(turn["momentum_x"][j]) + y.append(turn["position_y"][j]) + py.append(turn["momentum_y"][j]) + t.append(turn["position_t"][j]) + pt.append(turn["momentum_t"][j]) + n = n + 1 + +# Output number of periods in data series +nturns = len(x) +print(f"number of periods = {nturns}") +print() + +# Approximate the tune and closed orbit using the 4-turn formula: + +# from x data only +argument = (x[0] - x[1] + x[2] - x[3]) / (2.0 * (x[1] - x[2])) +tunex = np.arccos(argument) / (2.0 * np.pi) +print(f"tune output from 4-turn formula, using x data = {tunex}") + +# from y data only +argument = (y[0] - y[1] + y[2] - y[3]) / (2.0 * (y[1] - y[2])) +tuney = np.arccos(argument) / (2.0 * np.pi) +print(f"tune output from 4-turn formula, using y data = {tuney}") + +# from t data only +argument = (t[0] - t[1] + t[2] - t[3]) / (2.0 * (t[1] - t[2])) +tunet = np.arccos(argument) / (2.0 * np.pi) +print(f"tune output from 4-turn formula, using t data = {tunet}") + +rtol = 1.0e-3 +print(f" rtol={rtol}") + +assert np.allclose( + [tunex, tuney, tunet], + [ + 0.139, + 0.219, + 0.0250, + ], + rtol=rtol, +) diff --git a/examples/linear_map/input_map.in b/examples/linear_map/input_map.in new file mode 100644 index 000000000..737941289 --- /dev/null +++ b/examples/linear_map/input_map.in @@ -0,0 +1,57 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 45.6e3 +beam.charge = 2.72370027e-8 #population 1.7e11 +beam.particle = electron +beam.distribution = waterbag_from_twiss +beam.alphaX = 0.0 +beam.alphaY = 0.0 +beam.alphaT = 0.0 +beam.betaX = 0.15 +beam.betaY = 0.8e-3 +beam.betaT = 9.210526315789473 +beam.emittX = 0.27e-09 +beam.emittY = 1e-12 +beam.emittT = 1.33e-6 + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.periods = 1 +lattice.elements = monitor map1 + +monitor.type = beam_monitor +monitor.backend = h5 + +map1.type = linear_map +# horizontal plane +map1.R11 = 0.642252653176584 +map1.R12 = 0.114973951021402 +map1.R21 = -5.109953378728999 +map1.R22 = 0.642252653176584 +# vertical plane +map1.R33 = 0.193549468050860 +map1.R34 = 0.0007848724139547 +map1.R43 = -1226.363146804167548 +map1.R44 = 0.193549468050860 +# longitudinal plane +map1.R55 = 0.987688340595138 +map1.R56 = 1.440843756949495 +map1.R65 = -0.016984313347225 +map1.R66 = 0.987688340595138 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = false diff --git a/examples/linear_map/run_map.py b/examples/linear_map/run_map.py new file mode 100755 index 000000000..5268af4dd --- /dev/null +++ b/examples/linear_map/run_map.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements +#from elements import LinearTransport +import numpy as np + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 45.6e3 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 10000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + +# target beta functions (m) +beta_star_x = 0.15 +beta_star_y = 0.8e-3 +beta_star_t = 9.210526315789473 + +# particle bunch +distr = distribution.Waterbag( + **twiss( + beta_x=beta_star_x, + beta_y=beta_star_y, + beta_t=beta_star_t, + emitt_x=0.27e-09, + emitt_y=1.0e-12, + emitt_t=1.33e-06, + alpha_x=0.0, + alpha_y=0.0, + alpha_t=0.0, + ) +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# initialize the linear map +Iden = elements.LinearMap.Map6x6.identity() +Rmat = Iden + +# desired tunes +Qx = 0.139 +Qy = 0.219 +Qt = 0.0250 + +# desired phase advance +phi_x = 2.0*np.pi*Qx +phi_y = 2.0*np.pi*Qy +phi_t = 2.0*np.pi*Qt + +# matrix elements for the horizontal plane +R[1, 1] = np.cos(phi_x) +R[1, 2] = beta_star_x * np.sin(phi_x) +R[2, 1] = -np.sin(phi_x)/beta_star_x +R[2, 2] = np.cos(phi_x) +# matrix elements for the vertical plane +R[3, 3] = np.cos(phi_y) +R[3, 4] = beta_star_y * np.sin(phi_y) +R[4, 3] = -np.sin(phi_y)/beta_star_y +R[4, 4] = np.cos(phi_y) +# matrix elements for the longitudinal plane +R[5, 5] = np.cos(phi_t) +R[5, 6] = beta_star_t * np.sin(phi_t) +R[6, 5] = -np.sin(phi_t)/beta_star_t +R[6, 6] = np.cos(phi_t) + +# design the accelerator lattice +map = [ + monitor, + elements.LinearMap(R=Rmat), +] + +sim.lattice.extend(map) + +# number of periods through the lattice +sim.periods = 4 + +# run simulation +sim.evolve() + +# clean shutdown +sim.finalize() From 611a0ea1eaebbb81b8735cb307476a35041f619e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 12 Jan 2025 01:41:16 +0000 Subject: [PATCH 30/48] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/linear_map/analysis_map.py | 8 +++++--- examples/linear_map/run_map.py | 17 +++++++++-------- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/examples/linear_map/analysis_map.py b/examples/linear_map/analysis_map.py index 86cb5ec0c..4de79069f 100755 --- a/examples/linear_map/analysis_map.py +++ b/examples/linear_map/analysis_map.py @@ -9,6 +9,7 @@ import openpmd_api as io from scipy.stats import moment + def get_moments(beam): """Calculate standard deviations of beam position & momenta and emittance values @@ -31,6 +32,7 @@ def get_moments(beam): return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t) + # initial/final beam series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) last_step = list(series.iterations)[-1] @@ -48,7 +50,7 @@ def get_moments(beam): print( f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" ) - + atol = 0.0 # ignored rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") @@ -110,9 +112,9 @@ def get_moments(beam): turn = beam.to_df() x.append(turn["position_x"][j]) px.append(turn["momentum_x"][j]) - y.append(turn["position_y"][j]) + y.append(turn["position_y"][j]) py.append(turn["momentum_y"][j]) - t.append(turn["position_t"][j]) + t.append(turn["position_t"][j]) pt.append(turn["momentum_t"][j]) n = n + 1 diff --git a/examples/linear_map/run_map.py b/examples/linear_map/run_map.py index 5268af4dd..be0223c88 100755 --- a/examples/linear_map/run_map.py +++ b/examples/linear_map/run_map.py @@ -6,10 +6,11 @@ # # -*- coding: utf-8 -*- -from impactx import ImpactX, distribution, elements -#from elements import LinearTransport +# from elements import LinearTransport import numpy as np +from impactx import ImpactX, distribution, elements + sim = ImpactX() # set numerical parameters and IO control @@ -65,24 +66,24 @@ Qt = 0.0250 # desired phase advance -phi_x = 2.0*np.pi*Qx -phi_y = 2.0*np.pi*Qy -phi_t = 2.0*np.pi*Qt +phi_x = 2.0 * np.pi * Qx +phi_y = 2.0 * np.pi * Qy +phi_t = 2.0 * np.pi * Qt # matrix elements for the horizontal plane R[1, 1] = np.cos(phi_x) R[1, 2] = beta_star_x * np.sin(phi_x) -R[2, 1] = -np.sin(phi_x)/beta_star_x +R[2, 1] = -np.sin(phi_x) / beta_star_x R[2, 2] = np.cos(phi_x) # matrix elements for the vertical plane R[3, 3] = np.cos(phi_y) R[3, 4] = beta_star_y * np.sin(phi_y) -R[4, 3] = -np.sin(phi_y)/beta_star_y +R[4, 3] = -np.sin(phi_y) / beta_star_y R[4, 4] = np.cos(phi_y) # matrix elements for the longitudinal plane R[5, 5] = np.cos(phi_t) R[5, 6] = beta_star_t * np.sin(phi_t) -R[6, 5] = -np.sin(phi_t)/beta_star_t +R[6, 5] = -np.sin(phi_t) / beta_star_t R[6, 6] = np.cos(phi_t) # design the accelerator lattice From 70814163bafebfcc275fe315774d0cb5465a5ebf Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Sat, 11 Jan 2025 17:45:05 -0800 Subject: [PATCH 31/48] Document example README. --- docs/source/usage/examples.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 72ed8b3a6..8e911e9e6 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -35,7 +35,7 @@ Single Particle Dynamics examples/fodo_programmable/README.rst examples/dogleg/README.rst examples/coupled_optics/README.rst - + examples/linear_map/README.rst Collective Effects ------------------ From c81a0129e7e180137323668bfa5cf5c5b24b1639 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Sat, 11 Jan 2025 17:47:24 -0800 Subject: [PATCH 32/48] Update run_map.py Add twiss --- examples/linear_map/run_map.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/linear_map/run_map.py b/examples/linear_map/run_map.py index be0223c88..4496394e1 100755 --- a/examples/linear_map/run_map.py +++ b/examples/linear_map/run_map.py @@ -9,7 +9,7 @@ # from elements import LinearTransport import numpy as np -from impactx import ImpactX, distribution, elements +from impactx import ImpactX, distribution, elements, twiss sim = ImpactX() From 04be1fb349fc8a0d97729bf77bb215c664d8749e Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Sat, 11 Jan 2025 17:49:16 -0800 Subject: [PATCH 33/48] Update run_map.py Correct name of Rmat/R. --- examples/linear_map/run_map.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/linear_map/run_map.py b/examples/linear_map/run_map.py index 4496394e1..dd27301dd 100755 --- a/examples/linear_map/run_map.py +++ b/examples/linear_map/run_map.py @@ -58,7 +58,7 @@ # initialize the linear map Iden = elements.LinearMap.Map6x6.identity() -Rmat = Iden +R = Iden # desired tunes Qx = 0.139 From 1a36a89f2e52326cdc6aad68603ee7e1f149c806 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Sat, 11 Jan 2025 17:51:11 -0800 Subject: [PATCH 34/48] Update run_map.py Fix naming of R again. --- examples/linear_map/run_map.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/examples/linear_map/run_map.py b/examples/linear_map/run_map.py index dd27301dd..e65a77c2d 100755 --- a/examples/linear_map/run_map.py +++ b/examples/linear_map/run_map.py @@ -58,7 +58,7 @@ # initialize the linear map Iden = elements.LinearMap.Map6x6.identity() -R = Iden +Rmat = Iden # desired tunes Qx = 0.139 @@ -71,20 +71,20 @@ phi_t = 2.0 * np.pi * Qt # matrix elements for the horizontal plane -R[1, 1] = np.cos(phi_x) -R[1, 2] = beta_star_x * np.sin(phi_x) -R[2, 1] = -np.sin(phi_x) / beta_star_x -R[2, 2] = np.cos(phi_x) +Rmat[1, 1] = np.cos(phi_x) +Rmat[1, 2] = beta_star_x * np.sin(phi_x) +Rmat[2, 1] = -np.sin(phi_x) / beta_star_x +Rmat[2, 2] = np.cos(phi_x) # matrix elements for the vertical plane -R[3, 3] = np.cos(phi_y) -R[3, 4] = beta_star_y * np.sin(phi_y) -R[4, 3] = -np.sin(phi_y) / beta_star_y -R[4, 4] = np.cos(phi_y) +Rmat[3, 3] = np.cos(phi_y) +Rmat[3, 4] = beta_star_y * np.sin(phi_y) +Rmat[4, 3] = -np.sin(phi_y) / beta_star_y +Rmat[4, 4] = np.cos(phi_y) # matrix elements for the longitudinal plane -R[5, 5] = np.cos(phi_t) -R[5, 6] = beta_star_t * np.sin(phi_t) -R[6, 5] = -np.sin(phi_t) / beta_star_t -R[6, 6] = np.cos(phi_t) +Rmat[5, 5] = np.cos(phi_t) +Rmat[5, 6] = beta_star_t * np.sin(phi_t) +Rmat[6, 5] = -np.sin(phi_t) / beta_star_t +Rmat[6, 6] = np.cos(phi_t) # design the accelerator lattice map = [ From 4fa891bb19ba73e741fbff0493e6f474c2b1c539 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Sat, 11 Jan 2025 18:12:36 -0800 Subject: [PATCH 35/48] Update input_map.in Fix number of periods. --- examples/linear_map/input_map.in | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/linear_map/input_map.in b/examples/linear_map/input_map.in index 737941289..74dc06342 100644 --- a/examples/linear_map/input_map.in +++ b/examples/linear_map/input_map.in @@ -20,7 +20,7 @@ beam.emittT = 1.33e-6 ############################################################################### # Beamline: lattice elements and segments ############################################################################### -lattice.periods = 1 +lattice.periods = 5 lattice.elements = monitor map1 monitor.type = beam_monitor From 6c975f75751d416fdd201d49ed87514cb9e24acb Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sat, 11 Jan 2025 22:51:33 -0800 Subject: [PATCH 36/48] Apply suggestions from code review Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> --- docs/source/usage/workflows/add_element.rst | 4 ++-- examples/linear_map/README.rst | 26 ++++++++++----------- examples/linear_map/input_map.in | 1 + 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index 3610d8239..017e7da1c 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -23,8 +23,8 @@ The units of the transport matrix :math:`R` elements are... TODO .. note:: - Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. - Thus, where symplecticity is needed, the :ref:`Python Programmable Element ` and the :ref:`C++ Element ` provide a more general approach. + If a user-provided linear map is used, it is up to the user to ensure that the 6x6 transport matrix is symplectic. + If a more general form of user-defined transport is needed, the :ref:`Python Programmable Element ` and the :ref:`C++ Element ` provide a more general approach. .. _usage-workflows-add-element-python: diff --git a/examples/linear_map/README.rst b/examples/linear_map/README.rst index fad78a4ac..c6b72c746 100644 --- a/examples/linear_map/README.rst +++ b/examples/linear_map/README.rst @@ -1,7 +1,7 @@ .. _examples-linear-map: -Iteration of a user-defined linear map -======================================= +Iteration of a User-Defined Linear Map +====================================== This example illustrates the application of a user-defined linear map via a matrix. @@ -12,8 +12,8 @@ The (fractional) tunes (Qx, Qy, Qt) of the map are given by (0.139, 0.219, 0.025 We use a 45.6 GeV electron beam that is invariant under the action of the linear map (matched). The horizontal and vertical unnormalized emittances are 0.27 nm and 1.0 pm, respectively. -These parameters are based on the single-beam parameters of FCC-ee (Z-mode) appearing here: -https://twiki.cern.ch/twiki/bin/view/FCC/FCCeeParameters_CDRBaseline-1_0 +These parameters are based on the `single-beam parameters of FCC-ee (Z-mode) `__. +(`backup `__). The second moments of the phase space variables should be unchanged under application of the map. @@ -26,8 +26,8 @@ Run This example can be run **either** as: -* **Python** script: ``python3 run_multipole.py`` or -* ImpactX **executable** using an input file: ``impactx input_multipole.in`` +* **Python** script: ``python3 run_map.py`` or +* ImpactX **executable** using an input file: ``impactx input_map.in`` For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. @@ -35,15 +35,15 @@ For `MPI-parallel `__ runs, prefix these lines with ` .. tab-item:: Python: Script - .. literalinclude:: run_multipole.py + .. literalinclude:: run_map.py :language: python3 - :caption: You can copy this file from ``examples/multipole/run_multipole.py``. + :caption: You can copy this file from ``examples/linear_map/run_map.py``. .. tab-item:: Executable: Input File - .. literalinclude:: input_multipole.in + .. literalinclude:: input_map.in :language: ini - :caption: You can copy this file from ``examples/multipole/input_multipole.in``. + :caption: You can copy this file from ``examples/linear_map/input_map.in``. Analyze @@ -51,8 +51,8 @@ Analyze We run the following script to analyze correctness: -.. dropdown:: Script ``analysis_multipole.py`` +.. dropdown:: Script ``analysis_map.py`` - .. literalinclude:: analysis_multipole.py + .. literalinclude:: analysis_map.py :language: python3 - :caption: You can copy this file from ``examples/multipole/analysis_multipole.py``. + :caption: You can copy this file from ``examples/linear_map/analysis_map.py``. diff --git a/examples/linear_map/input_map.in b/examples/linear_map/input_map.in index 74dc06342..67fed5909 100644 --- a/examples/linear_map/input_map.in +++ b/examples/linear_map/input_map.in @@ -17,6 +17,7 @@ beam.emittX = 0.27e-09 beam.emittY = 1e-12 beam.emittT = 1.33e-6 + ############################################################################### # Beamline: lattice elements and segments ############################################################################### From 57c2f3d6911a50925df3e5773f08238053139036 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Sun, 12 Jan 2025 16:33:22 -0800 Subject: [PATCH 37/48] Add docs for matrix elements. --- docs/source/usage/parameters.rst | 9 +++++---- docs/source/usage/python.rst | 8 ++++++-- docs/source/usage/workflows/add_element.rst | 8 ++++++-- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 647f97a5a..5721b577b 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -351,15 +351,16 @@ Lattice Elements * ``linear_map`` for a custom, linear transport matrix. - The matrix elements are 1-indexed. - The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that are differing from that need to be specified. + The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. + The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. - TODO: describe units to put in the R entries. + The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m + and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. This requires these additional parameters: * ``.R(i,j)`` (``float``, ...) matrix entries - a 1-indexed, 6x6, linear transport map to multiply with the the beam vector :math:`x,px,y,py,t,pt`. + a 1-indexed, 6x6, linear transport map to multiply with the the phase space vector :math:`x,px,y,py,t,pt`. * ``.dx`` (``float``, in meters) horizontal translation error * ``.dy`` (``float``, in meters) vertical translation error * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index a6ebaa0c8..fc4f21407 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -615,9 +615,13 @@ This module provides elements for the accelerator lattice. A custom, linear transport matrix. - TODO: describe units to put in the R entries. + The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. + The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. - :param R: a linear transport map to multiply with the the beam vector :math:`x,px,y,py,t,pt`. + The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m + and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. + + :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`. :param dx: horizontal translation error in m :param dy: vertical translation error in m :param rotation: rotation error in the transverse plane [degrees] diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index 017e7da1c..ff09320e8 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -18,8 +18,12 @@ Linear Map A custom linear element can be provided by specifying the 6x6 linear transport matrix :math:`R` as an input. See the :ref:` example ` for Python and inputs file syntax to specify a custom linear element. -The entries of the transport matrix are ordered for multiplication with the beam vector :math:`x,px,y,py,t,pt`. -The units of the transport matrix :math:`R` elements are... TODO +The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. +The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. + +The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m +and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. + .. note:: From baf0fa02e21a36280b6378425802302778a013b7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jan 2025 00:33:50 +0000 Subject: [PATCH 38/48] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/source/usage/parameters.rst | 2 +- docs/source/usage/python.rst | 2 +- docs/source/usage/workflows/add_element.rst | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 5721b577b..336a2f197 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -351,7 +351,7 @@ Lattice Elements * ``linear_map`` for a custom, linear transport matrix. - The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. + The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index fc4f21407..c80b099b8 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -618,7 +618,7 @@ This module provides elements for the accelerator lattice. The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. - The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m + The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`. diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index ff09320e8..c801b2a25 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -21,7 +21,7 @@ See the :ref:` example ` for Python and inputs file synta The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. -The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m +The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. From d5897ac07492fd7e9dbf02293e45944224b6bddc Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 12 Jan 2025 22:50:27 -0800 Subject: [PATCH 39/48] Doc Updates --- docs/source/usage/python.rst | 1 - docs/source/usage/workflows/add_element.rst | 1 - 2 files changed, 2 deletions(-) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index c80b099b8..18ba2d169 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -616,7 +616,6 @@ This module provides elements for the accelerator lattice. A custom, linear transport matrix. The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. - The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index c801b2a25..fe8debaa3 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -19,7 +19,6 @@ A custom linear element can be provided by specifying the 6x6 linear transport m See the :ref:` example ` for Python and inputs file syntax to specify a custom linear element. The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. -The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. From aff0764a410fbb07abd47f3b895a8fb51c1a8953 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 13 Jan 2025 09:25:42 -0800 Subject: [PATCH 40/48] Generalize CMake Args for inputs Taken over from generalized version we merged to WarpX. --- examples/CMakeLists.txt | 37 +++++++++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index e9f6b29d1..e28590ff0 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -55,6 +55,30 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) # make a unique run directory file(MAKE_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}) + # get input file/script and optional command-line arguments + separate_arguments(INPUTS_LIST UNIX_COMMAND "${input}") + list(GET INPUTS_LIST 0 INPUTS_FILE) + list(LENGTH INPUTS_LIST INPUTS_LIST_LENGTH) + if(INPUTS_LIST_LENGTH GREATER 1) + list(SUBLIST INPUTS_LIST 1 -1 INPUTS_ARGS) + list(JOIN INPUTS_ARGS " " INPUTS_ARGS) + else() + set(INPUTS_ARGS "") + endif() + cmake_path(SET INPUTS_FILE "${ImpactX_SOURCE_DIR}/${INPUTS_FILE}") + + # get analysis script and optional command-line arguments + separate_arguments(ANALYSIS_LIST UNIX_COMMAND "${analysis_script}") + list(GET ANALYSIS_LIST 0 ANALYSIS_FILE) + cmake_path(SET ANALYSIS_FILE "${CMAKE_CURRENT_SOURCE_DIR}/${ANALYSIS_FILE}") + list(LENGTH ANALYSIS_LIST ANALYSIS_LIST_LENGTH) + if(ANALYSIS_LIST_LENGTH GREATER 1) + list(SUBLIST ANALYSIS_LIST 1 -1 ANALYSIS_ARGS) + list(JOIN ANALYSIS_ARGS " " ANALYSIS_ARGS) + else() + set(ANALYSIS_ARGS "") + endif() + # test run set(THIS_WORKING_DIR ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}) set(THIS_MPI_TEST_EXE) @@ -64,10 +88,10 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) set(THIS_Python_EXE) if(is_python) set(THIS_Python_EXE ${Python_EXECUTABLE}) - endif() - if(is_python) + # for argparse, do not pass command-line arguments as one quoted string + separate_arguments(INPUTS_ARGS UNIX_COMMAND "${INPUTS_ARGS}") add_test(NAME ${name}.run - COMMAND ${THIS_MPI_TEST_EXE} ${THIS_Python_EXE} ${ImpactX_SOURCE_DIR}/${input} + COMMAND ${THIS_MPI_TEST_EXE} ${THIS_Python_EXE} ${INPUTS_FILE} ${INPUTS_ARGS} WORKING_DIRECTORY ${THIS_WORKING_DIR} ) # TODO: @@ -77,12 +101,13 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) else() add_test(NAME ${name}.run COMMAND - ${THIS_MPI_TEST_EXE} $ ${ImpactX_SOURCE_DIR}/${input} + ${THIS_MPI_TEST_EXE} $ ${INPUTS_FILE} amrex.abort_on_unused_inputs=1 amrex.throw_exception = 1 amrex.signal_handling = 0 impactx.always_warn_immediately=1 - impactx.abort_on_warning_threshold=medium + impactx.abort_on_warning_threshold=low + ${INPUTS_ARGS} WORKING_DIRECTORY ${THIS_WORKING_DIR} ) endif() @@ -213,7 +238,7 @@ add_impactx_test(FODO.MADX.py # Python: FODO Cell w/ custom linear element ################################# # add_impactx_test(FODO.userdef - examples/fodo_userdef/input_fodo_userdef.in + "examples/fodo_userdef/input_fodo_userdef.in impactx.abort_on_warning_threshold=medium" ON # ImpactX MPI-parallel examples/fodo_userdef/analysis_fodo.py examples/fodo_userdef/plot_fodo.py From c8ac06c7083e46fe8ed90f802131c33eec87eb76 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 13 Jan 2025 09:45:38 -0800 Subject: [PATCH 41/48] Add symplectic warning in app element documentation. --- docs/source/usage/parameters.rst | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 336a2f197..4ce034d12 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -357,7 +357,9 @@ Lattice Elements The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. - This requires these additional parameters: + The internal tracking methods used by ImpactX are symplectic. However, if a user-defined linear map :math:`R` is provided, it is up to the user to ensure that the matrix :math:`R` is symplectic. Otherwise, this condition may be violated. + + This element requires these additional parameters: * ``.R(i,j)`` (``float``, ...) matrix entries a 1-indexed, 6x6, linear transport map to multiply with the the phase space vector :math:`x,px,y,py,t,pt`. From 46a0f5d48b261aa8ae4fda5dd01ddc0a4f65915f Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 13 Jan 2025 09:47:33 -0800 Subject: [PATCH 42/48] Add symplectic warning in Python element documentation. --- docs/source/usage/python.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 18ba2d169..0a7842fe2 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -620,6 +620,9 @@ This module provides elements for the accelerator lattice. The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. + The internal tracking methods used by ImpactX are symplectic. However, if a user-defined linear map :math:`R` is provided, it is + up to the user to ensure that the matrix :math:`R` is symplectic. Otherwise, this condition may be violated. + :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`. :param dx: horizontal translation error in m :param dy: vertical translation error in m From 697661b159c2ef0a2d074a027c4a08c03545b8ff Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jan 2025 17:47:42 +0000 Subject: [PATCH 43/48] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/source/usage/python.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 0a7842fe2..2f715df8b 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -620,7 +620,7 @@ This module provides elements for the accelerator lattice. The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. - The internal tracking methods used by ImpactX are symplectic. However, if a user-defined linear map :math:`R` is provided, it is + The internal tracking methods used by ImpactX are symplectic. However, if a user-defined linear map :math:`R` is provided, it is up to the user to ensure that the matrix :math:`R` is symplectic. Otherwise, this condition may be violated. :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`. From a1c2d808e21597a8442292d3e53597e221b25727 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 13 Jan 2025 10:19:48 -0800 Subject: [PATCH 44/48] Remove runtime warning. --- src/initialization/InitElement.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 67db6fcb3..1a7d55ef1 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -471,12 +471,6 @@ namespace detail pp_element.queryAddWithParser(name.c_str(), transport_map(i, j)); } } - ablastr::warn_manager::WMRecordWarning( - "ImpactX::read_element", - "Caution, a user-provided linear map is used in '" + - element_name + "'. Transport may not be symplectic.", - ablastr::warn_manager::WarnPriority::low - ); m_lattice.emplace_back(LinearMap(transport_map, a["dx"], a["dy"], a["rotation_degree"]) ); } else { From b81732b894c3df03876b31c6f2dd182bffe1c578 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Mon, 13 Jan 2025 12:16:46 -0800 Subject: [PATCH 45/48] Add support for nonzero ds to app input. --- docs/source/usage/parameters.rst | 1 + docs/source/usage/python.rst | 1 + examples/fodo_userdef/input_fodo_userdef.in | 2 ++ src/initialization/InitElement.cpp | 5 ++++- 4 files changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 4ce034d12..70d7437de 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -363,6 +363,7 @@ Lattice Elements * ``.R(i,j)`` (``float``, ...) matrix entries a 1-indexed, 6x6, linear transport map to multiply with the the phase space vector :math:`x,px,y,py,t,pt`. + * ``.ds`` (``float``, in meters) length associated with a user-defined linear element (defaults to 0) * ``.dx`` (``float``, in meters) horizontal translation error * ``.dy`` (``float``, in meters) vertical translation error * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 2f715df8b..889aca74d 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -624,6 +624,7 @@ This module provides elements for the accelerator lattice. up to the user to ensure that the matrix :math:`R` is symplectic. Otherwise, this condition may be violated. :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`. + :param ds: length associated with a user-defined linear element (defaults to 0), in m :param dx: horizontal translation error in m :param dy: vertical translation error in m :param rotation: rotation error in the transverse plane [degrees] diff --git a/examples/fodo_userdef/input_fodo_userdef.in b/examples/fodo_userdef/input_fodo_userdef.in index 0686b5fbf..420597477 100644 --- a/examples/fodo_userdef/input_fodo_userdef.in +++ b/examples/fodo_userdef/input_fodo_userdef.in @@ -28,6 +28,7 @@ monitor.type = beam_monitor monitor.backend = h5 drift1.type = linear_map +drift1.ds = 0.25 drift1.R12 = 0.25 # ds drift1.R34 = 0.25 # ds drift1.R56 = 0.25 / 16.6464 # ds / (beta*gamma^2) @@ -37,6 +38,7 @@ quad1.ds = 1.0 quad1.k = 1.0 drift2.type = linear_map +drift2.ds = 0.5 drift2.R12 = 0.5 # ds drift2.R34 = 0.5 # ds drift2.R56 = 0.5 / 16.6464 # ds / (beta*gamma^2) diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 1a7d55ef1..932ecdc3d 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -462,6 +462,9 @@ namespace detail { auto a = detail::query_alignment(pp_element); + amrex::ParticleReal ds = 0.0; + pp_element.queryAdd("ds", ds); + elements::LinearTransport::Map6x6 transport_map = elements::LinearTransport::Map6x6::Identity(); // safe to ParmParse inputs for reproducibility @@ -472,7 +475,7 @@ namespace detail } } - m_lattice.emplace_back(LinearMap(transport_map, a["dx"], a["dy"], a["rotation_degree"]) ); + m_lattice.emplace_back(LinearMap(transport_map, ds, a["dx"], a["dy"], a["rotation_degree"]) ); } else { amrex::Abort("Unknown type for lattice element " + element_name + ": " + element_type); } From bdb2d7f51dbb2cdf448ca62e439335881df0b7ac Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 13 Jan 2025 12:51:15 -0800 Subject: [PATCH 46/48] Update examples/CMakeLists.txt --- examples/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index e28590ff0..ce1704b09 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -238,7 +238,7 @@ add_impactx_test(FODO.MADX.py # Python: FODO Cell w/ custom linear element ################################# # add_impactx_test(FODO.userdef - "examples/fodo_userdef/input_fodo_userdef.in impactx.abort_on_warning_threshold=medium" + examples/fodo_userdef/input_fodo_userdef.in ON # ImpactX MPI-parallel examples/fodo_userdef/analysis_fodo.py examples/fodo_userdef/plot_fodo.py From d729b2f3b498b905383abde461121d8d25a233f8 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 13 Jan 2025 13:47:55 -0800 Subject: [PATCH 47/48] FODO Analysis: Check s and gamma --- examples/fodo/analysis_fodo.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/examples/fodo/analysis_fodo.py b/examples/fodo/analysis_fodo.py index f4e3e6b89..cb4b4dcdd 100755 --- a/examples/fodo/analysis_fodo.py +++ b/examples/fodo/analysis_fodo.py @@ -38,7 +38,8 @@ def get_moments(beam): series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) last_step = list(series.iterations)[-1] initial = series.iterations[1].particles["beam"].to_df() -final = series.iterations[last_step].particles["beam"].to_df() +beam_final = series.iterations[last_step].particles["beam"] +final = beam_final.to_df() # compare number of particles num_particles = 10000 @@ -74,9 +75,12 @@ def get_moments(beam): print("") print("Final Beam:") sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +s_ref = beam_final.get_attribute("s_ref") +gamma_ref = beam_final.get_attribute("gamma_ref") print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") print( - f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}\n" + f" s_ref={s_ref:e} gamma_ref={gamma_ref:e}" ) atol = 0.0 # ignored @@ -84,7 +88,7 @@ def get_moments(beam): print(f" rtol={rtol} (ignored: atol~={atol})") assert np.allclose( - [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref], [ 7.4790118496224206e-005, 7.5357525169680140e-005, @@ -92,6 +96,8 @@ def get_moments(beam): 1.9959539836392703e-009, 2.0175014668882125e-009, 2.0013820380883801e-006, + 3.000000, + 3.914902e003, ], rtol=rtol, atol=atol, From d0ffd9cf84a0f351da58235797751866baea2f6c Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 13 Jan 2025 13:53:17 -0800 Subject: [PATCH 48/48] __repr__: no select entries --- src/python/elements.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 256c9cf5f..ce712c370 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1578,9 +1578,7 @@ void init_elements(py::module& m) .def("__repr__", [](LinearMap const & linearmap) { return element_name( - linearmap, - std::make_pair("R11", linearmap.m_transport_map(1,1)), - std::make_pair("R12", linearmap.m_transport_map(1,2)) + linearmap ); } )