From 0be1a9bb422806cc52d38ca22d99962577782f76 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 9 Jan 2025 19:37:21 -0800 Subject: [PATCH] `LinearMap`: `ds` bookkeeping And reference particle pushing (drifting), if thick. --- src/particles/elements/LinearMap.H | 67 ++++++++++++++++++++++++++++-- src/python/elements.cpp | 15 ++++++- 2 files changed, 76 insertions(+), 6 deletions(-) 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);