diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 61cdd77aa..f11dc66ee 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,65 @@ namespace impactx { + /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix + */ + CovarianceMatrix + create_covariance_matrix ( + distribution::KnownDistributions const & distr + ) + { + // zero out the 6x6 matrix + CovarianceMatrix cv; + for (int i=1; i<=6; ++i) { + for (int j = 1; j <= 6; ++j) { + cv(i, j) = 0.0; + } + } + + // initialize from 2nd order beam moments + std::visit([&](auto&& distribution) { + // quick hack + using Distribution = std::remove_cv_t< std::remove_reference_t< decltype(distribution)> >; + if constexpr (std::is_same::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 = distribution.m_lambdaPx; + amrex::ParticleReal lambdaPy = distribution.m_lambdaPy; + amrex::ParticleReal lambdaPt = distribution.m_lambdaPt; + amrex::ParticleReal muxpx = distribution.m_muxpx; + amrex::ParticleReal muypy = distribution.m_muypy; + amrex::ParticleReal mutpt = distribution.m_mutpt; + + // use distribution inputs to populate a 6x6 covariance matrix + amrex::ParticleReal denom_x = 1.0 - muxpx*muxpx; + cv(1,1) = lambdaX*lambdaX / denom_x; + cv(1,2) = -lambdaX*lambdaPx*muxpx / denom_x; + cv(2,1) = cv(1,2); + cv(2,2) = lambdaPx*lambdaPx / denom_x; + + amrex::ParticleReal denom_y = 1.0 - muypy*muypy; + cv(3,3) = lambdaY*lambdaY / denom_y; + cv(3,4) = -lambdaY*lambdaPy*muypy / denom_y; + cv(4,3) = cv(3,4); + cv(4,4) = lambdaPy*lambdaPy / denom_y; + + amrex::ParticleReal denom_t = 1.0 - mutpt*mutpt; + cv(5,5) = lambdaT*lambdaT / denom_t; + cv(5,6) = -lambdaT*lambdaPt*mutpt / denom_t; + cv(6,5) = cv(5,6); + cv(6,6) = lambdaPt*lambdaPt / denom_t; + + } + }, distr); + + return cv; + } + void ImpactX::add_particles ( amrex::ParticleReal bunch_charge, @@ -95,7 +155,7 @@ namespace impactx amrex::ParticleReal * const AMREX_RESTRICT py_ptr = py.data(); amrex::ParticleReal * const AMREX_RESTRICT pt_ptr = pt.data(); - using Distribution = std::remove_reference_t< std::remove_cv_t >; + 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 71ec7e4aa..efbcd0390 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 @@ -436,6 +437,23 @@ namespace detail read_element(sub_element_name, m_lattice, nslice_default, mapsteps_default); } } + } else if (element_type == "linear_map") + { + auto a = detail::query_alignment(pp_element); + + elements::LinearTransport::Map6x6 transport_map; + for (int i=1; i<=6; ++i) { + for (int j=1; j<=6; ++j) { + amrex::ParticleReal R_ij = (i == j) ? 1.0 : 0.0; + std::string name = "R" + std::to_string(i) + std::to_string(j); + pp_element.queryAdd(name.c_str(), R_ij); + + transport_map(i, j) = R_ij; + } + } + std::cout << "Caution: A user-provided linear map is used. Transport may not be symplectic." << "\n"; + + m_lattice.emplace_back(LinearMap(transport_map, a["dx"], a["dy"], a["rotation_degree"]) ); } else { amrex::Abort("Unknown type for lattice element " + element_name + ": " + element_type); } diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H new file mode 100644 index 000000000..23af3e5aa --- /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::SmallMatrix; + +} // namespace impactx::distribution + +#endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H diff --git a/src/particles/PushAll.H b/src/particles/PushAll.H index 3fec0f3bc..453f1949d 100644 --- a/src/particles/PushAll.H +++ b/src/particles/PushAll.H @@ -51,6 +51,10 @@ namespace impactx element(ref_part); } + // push covariance matrix + // TODO + // note: decide what to do for elements that have no covariance matrix + // loop over refinement levels int const nLevel = pc.finestLevel(); for (int lev = 0; lev <= nLevel; ++lev) diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index ba108c9fe..27c194990 100644 --- a/src/particles/elements/All.H +++ b/src/particles/elements/All.H @@ -20,20 +20,21 @@ #include "ConstF.H" #include "DipEdge.H" #include "Drift.H" +#include "Empty.H" #include "ExactDrift.H" #include "ExactSbend.H" #include "Kicker.H" +#include "LinearMap.H" #include "Marker.H" #include "Multipole.H" -#include "Empty.H" #include "NonlinearLens.H" #include "Programmable.H" +#include "PRot.H" #include "Quad.H" #include "RFCavity.H" #include "Sbend.H" #include "ShortRF.H" #include "Sol.H" -#include "PRot.H" #include "SoftSol.H" #include "SoftQuad.H" #include "TaperedPL.H" @@ -61,6 +62,7 @@ namespace impactx ExactDrift, ExactSbend, Kicker, + LinearMap, Marker, Multipole, NonlinearLens, diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index fcb4ff6f8..944177558 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -16,9 +16,11 @@ #include "mixin/thick.H" #include "mixin/named.H" #include "mixin/nofinalize.H" +#include "mixin/lineartransport.H" #include #include +#include #include @@ -160,6 +162,40 @@ namespace impactx refpart.s = s + slice_ds; } + + /** This function returns the linear transport map. + * + * @returns 6x6 transport matrix + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + + amrex::SmallMatrix + transport_map(RefPart & AMREX_RESTRICT refpart) { + + using namespace amrex::literals; // for _rt and _prt + elements::LinearTransport::Map6x6 R = {}; + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // access reference particle values to find beta*gamma^2 + amrex::ParticleReal const pt_ref = refpart.pt; + amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + + // assign linear map matrix elements + R(1,1) = 1.0_prt; + R(1,2) = slice_ds; + R(2,2) = 1.0_prt; + R(3,3) = 1.0_prt; + R(3,4) = slice_ds; + R(4,4) = 1.0_prt; + R(5,5) = 1.0_prt; + R(5,6) = slice_ds/betgam2; + R(6,6) = 1.0_prt; + return R; + + } + }; } // namespace impactx diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H new file mode 100644 index 000000000..f50e765e9 --- /dev/null +++ b/src/particles/elements/LinearMap.H @@ -0,0 +1,135 @@ +/* Copyright 2022-2024 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ELEMENT_LINEAR_MAP_H +#define IMPACTX_ELEMENT_LINEAR_MAP_H + +#include "particles/ImpactXParticleContainer.H" +#include "mixin/alignment.H" +#include "mixin/beamoptic.H" +#include "mixin/lineartransport.H" +#include "mixin/thin.H" +#include "mixin/nofinalize.H" + +#include +#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 a user-provided linear transport map R + * to the 6-vector of phase space coordinates (x,px,y,py,t,pt). + * Thus x_final = R(1,1)*x + R(1,2)*px + R(1,3)*y + ..., + * px_final = R(2,1)*x + R(2,2)*px + R(2,3)*y + ..., etc. + * + * @param R user-provided transport map + * @param dx horizontal translation error in m + * @param dy vertical translation error in m + * @param rotation_degree rotation error in the transverse plane [degrees] + */ + LinearMap ( + LinearTransport::Map6x6 const & R, + amrex::ParticleReal dx = 0, + amrex::ParticleReal dy = 0, + amrex::ParticleReal rotation_degree = 0 + ) + : Alignment(dx, dy, rotation_degree) + { + for (int i=1; i<=6; ++i) { + for (int j = 1; j <= 6; ++j) { + m_transport_map(i, j) = R(i, j); + } + } + } + + /** Push all particles */ + using BeamOptic::operator(); + + /** This is a LinearMap functor, so that a variable of this type can be used like a + * LinearMap function. + * + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t (unused) + * @param px particle momentum in x + * @param py particle momentum in y + * @param pt particle momentum in t (unused) + * @param idcpu particle global index + * @param refpart reference particle (unused) + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + amrex::ParticleReal & AMREX_RESTRICT t, + amrex::ParticleReal & AMREX_RESTRICT px, + amrex::ParticleReal & AMREX_RESTRICT py, + amrex::ParticleReal & AMREX_RESTRICT pt, + [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu, + [[maybe_unused]] RefPart const & refpart + ) const + { + using namespace amrex::literals; // for _rt and _prt + + // shift due to alignment errors of the element + shift_in(x, y, px, py); + + // input and output phase space vectors + amrex::Array1D 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); + } + + /** This pushes the reference particle. */ + using Thin::operator(); + + LinearTransport::Map6x6 m_transport_map; // 6x6 transport map + + }; + +} // 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..46416ba90 --- /dev/null +++ b/src/particles/elements/mixin/lineartransport.H @@ -0,0 +1,51 @@ +/* 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 +#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::SmallMatrix; + Map6x6 m_transport_map; ///< linearized map + }; + +} // namespace impactx::elements + +#endif // IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H