From ba1d7815a8c25c60af7b127eb1f29dbc58cffdf6 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 19 Apr 2023 00:27:00 -0700 Subject: [PATCH] Update Particle Container to Pure SoA Transition particle containers to pure SoA layouts. --- cmake/dependencies/ABLASTR.cmake | 6 +- cmake/dependencies/pyAMReX.cmake | 2 +- examples/fodo/run_fodo_programmable.py | 20 ++-- src/initialization/InitMeshRefinement.cpp | 2 + src/initialization/InitParser.cpp | 5 + src/particles/ImpactXParticleContainer.H | 82 +++++++---------- src/particles/ImpactXParticleContainer.cpp | 42 ++++----- .../diagnostics/DiagnosticOutput.cpp | 47 +++++----- src/particles/elements/ConstF.H | 30 +++--- src/particles/elements/DipEdge.H | 14 ++- src/particles/elements/Drift.H | 30 +++--- src/particles/elements/Multipole.H | 30 +++--- src/particles/elements/None.H | 8 +- src/particles/elements/NonlinearLens.H | 30 +++--- src/particles/elements/PRot.H | 30 +++--- src/particles/elements/Quad.H | 36 ++++---- src/particles/elements/RFCavity.H | 29 +++--- src/particles/elements/Sbend.H | 30 +++--- src/particles/elements/ShortRF.H | 30 +++--- src/particles/elements/SoftQuad.H | 29 +++--- src/particles/elements/SoftSol.H | 29 +++--- src/particles/elements/Sol.H | 26 +++--- .../elements/diagnostics/openPMD.cpp | 92 ++++++++----------- src/particles/elements/mixin/beamoptic.H | 38 ++++---- src/particles/spacecharge/GatherAndPush.cpp | 16 ++-- .../CoordinateTransformation.cpp | 31 +++---- src/particles/transformation/ToFixedS.H | 20 ++-- src/particles/transformation/ToFixedT.H | 20 ++-- src/python/ImpactXParticleContainer.cpp | 18 ++-- 29 files changed, 415 insertions(+), 407 deletions(-) diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 134829094..34ff62617 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -166,10 +166,10 @@ set(ImpactX_openpmd_src "" "Local path to openPMD-api source directory (preferred if set)") # Git fetcher -set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git" +set(ImpactX_ablastr_repo "https://github.com/ax3l/WarpX.git" CACHE STRING "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)") -set(ImpactX_ablastr_branch "f71597fca2cf9d3700e7e11533f4a85a1846a2b8" +set(ImpactX_ablastr_branch "topic-particle-soa" CACHE STRING "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)") @@ -177,7 +177,7 @@ set(ImpactX_ablastr_branch "f71597fca2cf9d3700e7e11533f4a85a1846a2b8" set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git" CACHE STRING "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)") -set(ImpactX_amrex_branch "" +set(ImpactX_amrex_branch "ee492f47704ae8b094ad9bb4b6f758d82606a61e" CACHE STRING "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)") diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake index 7ddc80f8c..7f351e7c3 100644 --- a/cmake/dependencies/pyAMReX.cmake +++ b/cmake/dependencies/pyAMReX.cmake @@ -76,7 +76,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON) set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git" CACHE STRING "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)") -set(ImpactX_pyamrex_branch "a90fec615b314bb2d6143add9655cbf6afe61186" +set(ImpactX_pyamrex_branch "5e95674f150ed1e91f51eeb9383d079165da9255" CACHE STRING "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)") diff --git a/examples/fodo/run_fodo_programmable.py b/examples/fodo/run_fodo_programmable.py index 1e1edf0e5..11c2cdd49 100755 --- a/examples/fodo/run_fodo_programmable.py +++ b/examples/fodo/run_fodo_programmable.py @@ -78,16 +78,16 @@ def my_drift(pge, pti, refpart): else: array = np.array - # access AoS data such as positions and cpu/id - aos = pti.aos() - aos_arr = array(aos, copy=False) - # access SoA data such as momentum + # access particle attributes soa = pti.soa() real_arrays = soa.GetRealData() - px = array(real_arrays[0], copy=False) - py = array(real_arrays[1], copy=False) - pt = array(real_arrays[2], copy=False) + x = array(real_arrays[0], copy=False) + y = array(real_arrays[1], copy=False) + t = array(real_arrays[2], copy=False) + px = array(real_arrays[3], copy=False) + py = array(real_arrays[4], copy=False) + pt = array(real_arrays[5], copy=False) # length of the current slice slice_ds = pge.ds / pge.nslice @@ -97,9 +97,9 @@ def my_drift(pge, pti, refpart): betgam2 = pt_ref**2 - 1.0 # advance position and momentum (drift) - aos_arr[:]["x"] += slice_ds * px[:] - aos_arr[:]["y"] += slice_ds * py[:] - aos_arr[:]["z"] += (slice_ds / betgam2) * pt[:] + x[:] += slice_ds * px[:] + y[:] += slice_ds * py[:] + t[:] += (slice_ds / betgam2) * pt[:] def my_ref_drift(pge, refpart): diff --git a/src/initialization/InitMeshRefinement.cpp b/src/initialization/InitMeshRefinement.cpp index 9769d4ee1..90eb19ee0 100644 --- a/src/initialization/InitMeshRefinement.cpp +++ b/src/initialization/InitMeshRefinement.cpp @@ -201,6 +201,8 @@ namespace impactx amrex::Geometry g = Geom(lev); g.ProbDomain(rb); amrex::AmrMesh::SetGeometry(lev, g); + + amrex::Print() << "New Geometry g=" << g << "\n"; } } } // namespace impactx diff --git a/src/initialization/InitParser.cpp b/src/initialization/InitParser.cpp index e6aa560bf..d16507037 100644 --- a/src/initialization/InitParser.cpp +++ b/src/initialization/InitParser.cpp @@ -25,6 +25,11 @@ namespace impactx::initialization bool the_arena_is_managed = false; // AMReX' default: true pp_amrex.queryAdd("the_arena_is_managed", the_arena_is_managed); + // debugging + // https://amrex-codes.github.io/amrex/docs_html/Debugging.html#breaking-into-debuggers + pp_amrex.add("throw_exception", 1); + pp_amrex.add("signal_handling", 0); + // Here we override the default tiling option for particles, which is always // "false" in AMReX, to "false" if compiling for GPU execution and "true" // if compiling for CPU. diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H index dd8304115..12ebe1fe6 100644 --- a/src/particles/ImpactXParticleContainer.H +++ b/src/particles/ImpactXParticleContainer.H @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -28,43 +29,16 @@ namespace impactx { - /** AMReX pre-defined Real attributes - * - * These are the AMReX pre-defined struct indexes for the Real attributes - * stored in an AoS in ImpactXParticleContainer. We document this here, - * because we change the meaning of these "positions" depending on the - * coordinate system we are currently in. - */ - struct RealAoS - { - enum - { - x, ///< position in x [m] (at fixed s OR fixed t) - y, ///< position in y [m] (at fixed s OR fixed t) - t, ///< time-of-flight c*t [m] (at fixed s) - nattribs ///< the number of attributes above (always last) - }; - - // at fixed t, the third component represents the position z - enum { - z = t ///< position in z [m] (at fixed t) - }; - - //! named labels for fixed s - static constexpr auto names_s = { "position_x", "position_y", "position_t" }; - //! named labels for fixed t - static constexpr auto names_t = { "position_x", "position_y", "position_z" }; - static_assert(names_s.size() == nattribs); - static_assert(names_t.size() == nattribs); - }; - - /** This struct indexes the additional Real attributes + /** This struct indexes the Real attributes * stored in an SoA in ImpactXParticleContainer */ struct RealSoA { enum { + x, ///< position in x [m] (at fixed s or t) + y, ///< position in y [m] (at fixed s or t) + t, ///< time-of-flight ct [m] (at fixed s) ux, ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t) uy, ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t) pt, ///< energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s) @@ -73,28 +47,38 @@ namespace impactx nattribs ///< the number of attributes above (always last) }; - // at fixed t, the third component represents the momentum in z + // at fixed t, the third component represents the position z, the 6th component represents the momentum in z enum { + z = t, ///< position in z [m] (at fixed t) pz = pt ///< momentum in z, scaled by the magnitude of the reference momentum [unitless] (at fixed t) }; //! named labels for fixed s - static constexpr auto names_s = { "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" }; + static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" }; //! named labels for fixed t - static constexpr auto names_t = { "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" }; + static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" }; static_assert(names_s.size() == nattribs); static_assert(names_t.size() == nattribs); }; - /** This struct indexes the additional Integer attributes + /** This struct indexes the Integer attributes * stored in an SoA in ImpactXParticleContainer */ struct IntSoA { enum { - nattribs ///< the number of particles above (always last) + id, ///< unique particle id on the rank of particle creation + cpu, ///< rank the particle was created on + nattribs ///< the number of attributes above (always last) }; + + //! named labels for fixed s + static constexpr auto names_s = { "id", "cpu" }; + //! named labels for fixed t + static constexpr auto names_t = { "id", "cpu" }; + static_assert(names_s.size() == nattribs); + static_assert(names_t.size() == nattribs); }; /** AMReX iterator for particle boxes @@ -102,15 +86,15 @@ namespace impactx * We subclass here to change the default threading strategy, which is * `static` in AMReX, to `dynamic` in ImpactX. */ - class ParIter - : public amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> + class ParIterSoA + : public amrex::ParIterSoA { public: - using amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParIter; + using amrex::ParIterSoA::ParIterSoA; - ParIter (ContainerType& pc, int level); + ParIterSoA (ContainerType& pc, int level); - ParIter (ContainerType& pc, int level, amrex::MFItInfo& info); + ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info); }; /** Const AMReX iterator for particle boxes - data is read only. @@ -118,15 +102,15 @@ namespace impactx * We subclass here to change the default threading strategy, which is * `static` in AMReX, to `dynamic` in ImpactX. */ - class ParConstIter - : public amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> + class ParConstIterSoA + : public amrex::ParConstIterSoA { public: - using amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParConstIter; + using amrex::ParConstIterSoA::ParConstIterSoA; - ParConstIter (ContainerType& pc, int level); + ParConstIterSoA (ContainerType& pc, int level); - ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info); + ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info); }; /** Beam Particles in ImpactX @@ -134,14 +118,14 @@ namespace impactx * This class stores particles, distributed over MPI ranks. */ class ImpactXParticleContainer - : public amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs> + : public amrex::ParticleContainerPureSoA { public: //! amrex iterator for particle boxes - using iterator = impactx::ParIter; + using iterator = impactx::ParIterSoA; //! amrex constant iterator for particle boxes (read-only) - using const_iterator = impactx::ParConstIter; + using const_iterator = impactx::ParConstIterSoA; //! Construct a new particle container ImpactXParticleContainer (amrex::AmrCore* amr_core); diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index 165b1295b..f91b2a812 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -17,7 +17,6 @@ #include #include #include -#include #include @@ -31,24 +30,24 @@ namespace impactx return do_dynamic; } - ParIter::ParIter (ContainerType& pc, int level) - : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParIterSoA::ParIterSoA (ContainerType& pc, int level) + : amrex::ParIterSoA(pc, level, amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {} - ParIter::ParIter (ContainerType& pc, int level, amrex::MFItInfo& info) - : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParIterSoA::ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info) + : amrex::ParIterSoA(pc, level, info.SetDynamic(do_omp_dynamic())) {} - ParConstIter::ParConstIter (ContainerType& pc, int level) - : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level) + : amrex::ParConstIterSoA(pc, level, amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {} - ParConstIter::ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info) - : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info) + : amrex::ParConstIterSoA(pc, level, info.SetDynamic(do_omp_dynamic())) {} ImpactXParticleContainer::ImpactXParticleContainer (amrex::AmrCore* amr_core) - : amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB()) + : amrex::ParticleContainerPureSoA(amr_core->GetParGDB()) { SetParticleSize(); } @@ -105,35 +104,26 @@ namespace impactx reserveData(); resizeData(); + // write Real attributes (SoA) to particle initialized zero auto& particle_tile = DefineAndReturnParticleTile(0, 0, 0); /* Create a temporary tile to obtain data from simulation. This data * is then copied to the permanent tile which is stored on the particle * (particle_tile). */ - using PinnedTile = amrex::ParticleTile< - amrex::Particle, - NArrayReal, NArrayInt, - amrex::PinnedArenaAllocator - >; + using PinnedTile = typename ContainerLike::ParticleTileType; PinnedTile pinned_tile; pinned_tile.define(NumRuntimeRealComps(), NumRuntimeIntComps()); for (int i = 0; i < np; i++) { - ParticleType p; - p.id() = ParticleType::NextID(); - p.cpu() = amrex::ParallelDescriptor::MyProc(); - p.pos(RealAoS::x) = x[i]; - p.pos(RealAoS::y) = y[i]; - p.pos(RealAoS::t) = t[i]; - // write position, creating cpu id, and particle id - pinned_tile.push_back(p); + pinned_tile.push_back_int(IntSoA::id, ParticleType::NextID()); + pinned_tile.push_back_int(IntSoA::cpu, amrex::ParallelDescriptor::MyProc()); } - // write Real attributes (SoA) to particle initialized zero - DefineAndReturnParticleTile(0, 0, 0); - + pinned_tile.push_back_real(RealSoA::x, x); + pinned_tile.push_back_real(RealSoA::y, y); + pinned_tile.push_back_real(RealSoA::t, t); pinned_tile.push_back_real(RealSoA::ux, px); pinned_tile.push_back_real(RealSoA::uy, py); pinned_tile.push_back_real(RealSoA::pt, pt); diff --git a/src/particles/diagnostics/DiagnosticOutput.cpp b/src/particles/diagnostics/DiagnosticOutput.cpp index 538e1b498..70abbe1bb 100644 --- a/src/particles/diagnostics/DiagnosticOutput.cpp +++ b/src/particles/diagnostics/DiagnosticOutput.cpp @@ -17,6 +17,7 @@ #include // for ParmParse #include // for ParticleReal #include // for PrintToFile +#include // for constructor of SoAParticle namespace impactx::diagnostics @@ -46,7 +47,8 @@ namespace impactx::diagnostics } // create a host-side particle buffer - auto tmp = pc.make_alike(); + using ParticleType = amrex::SoAParticle; + auto tmp = pc.template make_alike(); // copy device-to-host bool const local = true; @@ -56,34 +58,32 @@ namespace impactx::diagnostics int const nLevel = tmp.finestLevel(); for (int lev = 0; lev <= nLevel; ++lev) { // loop over all particle boxes - using ParIt = typename decltype(tmp)::ParConstIterType; - for (ParIt pti(tmp, lev); pti.isValid(); ++pti) { + using MyPinnedParIter = amrex::ParIterSoA; + for (MyPinnedParIter pti(tmp, lev); pti.isValid(); ++pti) { const int np = pti.numParticles(); - // preparing access to particle data: AoS - using PType = ImpactXParticleContainer::ParticleType; - auto const &aos = pti.GetArrayOfStructs(); - PType const *const AMREX_RESTRICT aos_ptr = aos().dataPtr(); - // preparing access to particle data: SoA of Reals - auto &soa_real = pti.GetStructOfArrays().GetRealData(); - amrex::ParticleReal const *const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr(); - amrex::ParticleReal const *const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr(); - amrex::ParticleReal const *const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr(); + auto const& particle_attributes = pti.GetStructOfArrays(); + auto const& part_x = particle_attributes.GetRealData(RealSoA::x); + auto const& part_y = particle_attributes.GetRealData(RealSoA::y); + auto const& part_t = particle_attributes.GetRealData(RealSoA::t); + auto const& part_px = particle_attributes.GetRealData(RealSoA::ux); + auto const& part_py = particle_attributes.GetRealData(RealSoA::uy); + auto const& part_pt = particle_attributes.GetRealData(RealSoA::pt); + + // Preparing the constructor for the new SoA Particle Type + auto ptd = pti.GetParticleTile().getParticleTileData(); if (otype == OutputType::PrintParticles) { // print out particles (this hack works only on CPU and on GPUs with // unified memory access) for (int i = 0; i < np; ++i) { - - // access AoS data such as positions and cpu/id - PType const &p = aos_ptr[i]; - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); + ParticleType p(ptd, i); uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu()); - // access SoA Real data + amrex::ParticleReal const x = part_x[i]; + amrex::ParticleReal const y = part_y[i]; + amrex::ParticleReal const t = part_t[i]; amrex::ParticleReal const px = part_px[i]; amrex::ParticleReal const py = part_py[i]; amrex::ParticleReal const pt = part_pt[i]; @@ -119,14 +119,11 @@ namespace impactx::diagnostics // print out particles (this hack works only on CPU and on GPUs with // unified memory access) for (int i = 0; i < np; ++i) { - - // access AoS data such as positions and cpu/id - PType const &p = aos_ptr[i]; - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); + ParticleType p(ptd, i); + amrex::ParticleReal const x = p.pos(RealSoA::x); + amrex::ParticleReal const y = p.pos(RealSoA::y); uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu()); - // access SoA Real data amrex::ParticleReal const px = part_px[i]; amrex::ParticleReal const py = part_py[i]; diff --git a/src/particles/elements/ConstF.H b/src/particles/elements/ConstF.H index 98ccdf638..824fecbc9 100644 --- a/src/particles/elements/ConstF.H +++ b/src/particles/elements/ConstF.H @@ -51,7 +51,9 @@ namespace impactx /** This pushes a single particle, relative to the reference particle * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -59,7 +61,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -67,16 +71,14 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // access reference particle values to find beta*gamma^2 amrex::ParticleReal const pt_ref = refpart.pt; amrex::ParticleReal const betgam2 = pow(pt_ref, 2) - 1.0_prt; - // intialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -85,20 +87,22 @@ namespace impactx amrex::ParticleReal const slice_ds = m_ds / nslice(); // advance position and momentum - p.pos(RealAoS::x) = cos(m_kx*slice_ds)*x + sin(m_kx*slice_ds)/m_kx*px; + xout = cos(m_kx*slice_ds)*x + sin(m_kx*slice_ds)/m_kx*px; pxout = -m_kx*sin(m_kx*slice_ds)*x + cos(m_kx*slice_ds)*px; - p.pos(RealAoS::y) = cos(m_ky*slice_ds)*y + sin(m_ky*slice_ds)/m_ky*py; + yout = cos(m_ky*slice_ds)*y + sin(m_ky*slice_ds)/m_ky*py; pyout = -m_ky*sin(m_ky*slice_ds)*y + cos(m_ky*slice_ds)*py; - p.pos(RealAoS::t) = cos(m_kt*slice_ds)*t + sin(m_kt*slice_ds)/(betgam2*m_kt)*pt; + tout = cos(m_kt*slice_ds)*t + sin(m_kt*slice_ds)/(betgam2*m_kt)*pt; ptout = -(m_kt*betgam2)*sin(m_kt*slice_ds)*t + cos(m_kt*slice_ds)*pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. diff --git a/src/particles/elements/DipEdge.H b/src/particles/elements/DipEdge.H index c86b2ee68..4846fc4e3 100644 --- a/src/particles/elements/DipEdge.H +++ b/src/particles/elements/DipEdge.H @@ -57,7 +57,9 @@ namespace impactx /** This is a dipedge functor, so that a variable of this type can be used like a * dipedge function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t (unchanged) @@ -65,7 +67,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT t, amrex::ParticleReal & AMREX_RESTRICT px, amrex::ParticleReal & AMREX_RESTRICT py, [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT pt, @@ -73,12 +77,6 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - - // access reference particle values if needed - // edge focusing matrix elements (zero gap) amrex::ParticleReal const R21 = tan(m_psi)/m_rc; amrex::ParticleReal R43 = -R21; diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index f14ba3fb9..cfdc0cc8f 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -46,7 +46,9 @@ namespace impactx /** This is a drift functor, so that a variable of this type can be used like a drift function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -54,7 +56,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -62,12 +66,10 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - - // initialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -80,18 +82,20 @@ namespace impactx amrex::ParticleReal const betgam2 = pow(pt_ref, 2) - 1.0_prt; // advance position and momentum (drift) - p.pos(RealAoS::x) = x + slice_ds * px; + xout = x + slice_ds * px; // pxout = px; - p.pos(RealAoS::y) = y + slice_ds * py; + yout = y + slice_ds * py; // pyout = py; - p.pos(RealAoS::t) = t + (slice_ds/betgam2) * pt; + tout = t + (slice_ds/betgam2) * pt; // ptout = pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. diff --git a/src/particles/elements/Multipole.H b/src/particles/elements/Multipole.H index ef487a8c1..ef5b1606a 100644 --- a/src/particles/elements/Multipole.H +++ b/src/particles/elements/Multipole.H @@ -56,7 +56,9 @@ namespace impactx /** This is a multipole functor, so that a variable of this type can be used like a * multipole function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -64,7 +66,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT t, amrex::ParticleReal & AMREX_RESTRICT px, amrex::ParticleReal & AMREX_RESTRICT py, amrex::ParticleReal & AMREX_RESTRICT pt, @@ -75,16 +79,14 @@ namespace impactx // a complex type with two amrex::ParticleReal using Complex = amrex::GpuComplex; - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // access reference particle values to find (beta*gamma)^2 //amrex::ParticleReal const pt_ref = refpart.pt; //amrex::ParticleReal const betgam2 = pow(pt_ref, 2) - 1.0_prt; - // intialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -101,20 +103,22 @@ namespace impactx amrex::ParticleReal const dpy = kick.m_imag/m_mfactorial; // advance position and momentum - p.pos(RealAoS::x) = x; + // xout = x; pxout = px + dpx; - p.pos(RealAoS::y) = y; + // yout = y; pyout = py + dpy; - p.pos(RealAoS::t) = t; + // tout = t; ptout = pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. */ diff --git a/src/particles/elements/None.H b/src/particles/elements/None.H index ecad8223f..0c3358cb3 100644 --- a/src/particles/elements/None.H +++ b/src/particles/elements/None.H @@ -51,7 +51,9 @@ namespace impactx /** Does nothing to a particle. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -59,7 +61,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - [[maybe_unused]] PType& AMREX_RESTRICT p, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT x, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT y, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT t, [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT px, [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT py, [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT pt, diff --git a/src/particles/elements/NonlinearLens.H b/src/particles/elements/NonlinearLens.H index 42f6431ff..1a9e7377a 100644 --- a/src/particles/elements/NonlinearLens.H +++ b/src/particles/elements/NonlinearLens.H @@ -53,7 +53,9 @@ namespace impactx /** This is a nonlinear lens functor, so that a variable of this type * can be used like a nonlinear lens function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -61,7 +63,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT t, amrex::ParticleReal & AMREX_RESTRICT px, amrex::ParticleReal & AMREX_RESTRICT py, amrex::ParticleReal & AMREX_RESTRICT pt, @@ -72,16 +76,14 @@ namespace impactx // a complex type with two amrex::ParticleReal using Complex = amrex::GpuComplex; - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // access reference particle values to find (beta*gamma)^2 //amrex::ParticleReal const pt_ref = refpart.pt; //amrex::ParticleReal const betgam2 = pow(pt_ref, 2) - 1.0_prt; - // intialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -110,20 +112,22 @@ namespace impactx amrex::ParticleReal dpy = -kick*dF.m_imag; // advance position and momentum - p.pos(RealAoS::x) = x; + // xout = x; pxout = px + dpx; - p.pos(RealAoS::y) = y; + // yout = y; pyout = py + dpy; - p.pos(RealAoS::t) = t; + // tout = t; ptout = pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. */ diff --git a/src/particles/elements/PRot.H b/src/particles/elements/PRot.H index 87a712e7e..1977ac6ea 100644 --- a/src/particles/elements/PRot.H +++ b/src/particles/elements/PRot.H @@ -54,7 +54,9 @@ namespace impactx /** This is a prot functor, so that a variable of this type can be used like a * prot function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -62,7 +64,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -70,15 +74,13 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // access reference particle values to find beta: amrex::ParticleReal const beta = refpart.beta(); - // intialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -91,20 +93,22 @@ namespace impactx sin(m_phi_in))*sin(theta); // advance position and momentum - p.pos(RealAoS::x) = x*pz/pzf; + xout = x*pz/pzf; pxout = px*cos(theta) + (pz - cos(m_phi_in))*sin(theta); - p.pos(RealAoS::y) = y + py*x*sin(theta)/pzf; + yout = y + py*x*sin(theta)/pzf; pyout = py; - p.pos(RealAoS::t) = t - (pt - 1.0_prt/beta)*x*sin(theta)/pzf; + tout = t - (pt - 1.0_prt/beta)*x*sin(theta)/pzf; ptout = pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. */ diff --git a/src/particles/elements/Quad.H b/src/particles/elements/Quad.H index 30fd2057b..24cc4bcde 100644 --- a/src/particles/elements/Quad.H +++ b/src/particles/elements/Quad.H @@ -51,7 +51,9 @@ namespace impactx /** This is a quad functor, so that a variable of this type can be used like a quad function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -59,7 +61,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -67,11 +71,6 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // length of the current slice amrex::ParticleReal const slice_ds = m_ds / nslice(); @@ -82,38 +81,43 @@ namespace impactx // compute phase advance per unit length in s (in rad/m) amrex::ParticleReal const omega = sqrt(std::abs(m_k)); - // intialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; if(m_k > 0.0) { // advance position and momentum (focusing quad) - p.pos(RealAoS::x) = cos(omega*slice_ds)*x + sin(omega*slice_ds)/omega*px; + xout = cos(omega*slice_ds)*x + sin(omega*slice_ds)/omega*px; pxout = -omega*sin(omega*slice_ds)*x + cos(omega*slice_ds)*px; - p.pos(RealAoS::y) = cosh(omega*slice_ds)*y + sinh(omega*slice_ds)/omega*py; + yout = cosh(omega*slice_ds)*y + sinh(omega*slice_ds)/omega*py; pyout = omega*sinh(omega*slice_ds)*y + cosh(omega*slice_ds)*py; - p.pos(RealAoS::t) = t + (slice_ds/betgam2)*pt; + tout = t + (slice_ds/betgam2)*pt; // ptout = pt; } else { // advance position and momentum (defocusing quad) - p.pos(RealAoS::x) = cosh(omega*slice_ds)*x + sinh(omega*slice_ds)/omega*px; + xout = cosh(omega*slice_ds)*x + sinh(omega*slice_ds)/omega*px; pxout = omega*sinh(omega*slice_ds)*x + cosh(omega*slice_ds)*px; - p.pos(RealAoS::y) = cos(omega*slice_ds)*y + sin(omega*slice_ds)/omega*py; + yout = cos(omega*slice_ds)*y + sin(omega*slice_ds)/omega*py; pyout = -omega*sin(omega*slice_ds)*y + cos(omega*slice_ds)*py; - p.pos(RealAoS::t) = t + (slice_ds/betgam2)*pt; + tout = t + (slice_ds/betgam2)*pt; // ptout = pt; } - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. diff --git a/src/particles/elements/RFCavity.H b/src/particles/elements/RFCavity.H index ce6b7d878..f652f4843 100644 --- a/src/particles/elements/RFCavity.H +++ b/src/particles/elements/RFCavity.H @@ -168,7 +168,9 @@ namespace RFCavityData /** This is an RF cavity functor, so that a variable of this type can be used like * an RF cavity function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -176,7 +178,9 @@ namespace RFCavityData */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -185,12 +189,10 @@ namespace RFCavityData { using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - - // initialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -205,20 +207,23 @@ namespace RFCavityData // so that, e.g., R(3,4) = dyf/dpyi. // push particles using the linear map - p.pos(RealAoS::x) = R(1,1)*x + R(1,2)*px + R(1,3)*y + xout = R(1,1)*x + R(1,2)*px + R(1,3)*y + R(1,4)*py + R(1,5)*t + R(1,6)*pt; pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y + R(2,4)*py + R(2,5)*t + R(2,6)*pt; - p.pos(RealAoS::y) = R(3,1)*x + R(3,2)*px + R(3,3)*y + yout = R(3,1)*x + R(3,2)*px + R(3,3)*y + R(3,4)*py + R(3,5)*t + R(3,6)*pt; pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y + R(4,4)*py + R(4,5)*t + R(4,6)*pt; - p.pos(RealAoS::t) = R(5,1)*x + R(5,2)*px + R(5,3)*y + tout = R(5,1)*x + R(5,2)*px + R(5,3)*y + R(5,4)*py + R(5,5)*t + R(5,6)*pt; ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y + R(6,4)*py + R(6,5)*t + R(6,6)*pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; diff --git a/src/particles/elements/Sbend.H b/src/particles/elements/Sbend.H index ac7d49e3a..c4d63c29f 100644 --- a/src/particles/elements/Sbend.H +++ b/src/particles/elements/Sbend.H @@ -48,7 +48,9 @@ namespace impactx /** This is a sbend functor, so that a variable of this type can be used like a sbend function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -56,7 +58,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -64,12 +68,10 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - - // initialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -89,25 +91,27 @@ namespace impactx amrex::ParticleReal const cos_theta = cos(theta); // advance position and momentum (sector bend) - p.pos(RealAoS::x) = cos_theta*x + m_rc*sin_theta*px + xout = cos_theta*x + m_rc*sin_theta*px - (m_rc/bet)*(1.0_prt - cos_theta)*pt; pxout = -sin_theta/m_rc*x + cos_theta*px - sin_theta/bet*pt; - p.pos(RealAoS::y) = y + m_rc*theta*py; + yout = y + m_rc*theta*py; // pyout = py; - p.pos(RealAoS::t) = sin_theta/bet*x + m_rc/bet*(1.0_prt - cos_theta)*px + t + tout = sin_theta/bet*x + m_rc/bet*(1.0_prt - cos_theta)*px + t + m_rc*(-theta+sin_theta/(bet*bet))*pt; // ptout = pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. diff --git a/src/particles/elements/ShortRF.H b/src/particles/elements/ShortRF.H index a3ff4645e..9f5dc1e2d 100644 --- a/src/particles/elements/ShortRF.H +++ b/src/particles/elements/ShortRF.H @@ -47,7 +47,9 @@ namespace impactx /** This is a shortrf functor, so that a variable of this type can be used like a * shortrf function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -55,7 +57,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -63,35 +67,35 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // access reference particle values to find (beta*gamma)^2 amrex::ParticleReal const pt_ref = refpart.pt; amrex::ParticleReal const betgam2 = pow(pt_ref, 2) - 1.0_prt; - // intialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; // advance position and momentum - p.pos(RealAoS::x) = x; + // xout = x; pxout = px + m_k*m_V/(2.0_prt*betgam2)*x; - p.pos(RealAoS::y) = y; + // yout = y; pyout = py + m_k*m_V/(2.0_prt*betgam2)*y; - p.pos(RealAoS::t) = t; + // tout = t; ptout = pt - m_k*m_V*t; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. */ diff --git a/src/particles/elements/SoftQuad.H b/src/particles/elements/SoftQuad.H index 885b27af0..ad55174ba 100644 --- a/src/particles/elements/SoftQuad.H +++ b/src/particles/elements/SoftQuad.H @@ -173,7 +173,9 @@ namespace SoftQuadrupoleData /** This is a soft-edge quadrupole functor, so that a variable of this type can be used * like a soft-edge quadrupole function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -181,7 +183,9 @@ namespace SoftQuadrupoleData */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -190,12 +194,10 @@ namespace SoftQuadrupoleData { using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - - // initialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -210,20 +212,23 @@ namespace SoftQuadrupoleData // so that, e.g., R(3,4) = dyf/dpyi. // push particles using the linear map - p.pos(RealAoS::x) = R(1,1)*x + R(1,2)*px + R(1,3)*y + xout = R(1,1)*x + R(1,2)*px + R(1,3)*y + R(1,4)*py + R(1,5)*t + R(1,6)*pt; pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y + R(2,4)*py + R(2,5)*t + R(2,6)*pt; - p.pos(RealAoS::y) = R(3,1)*x + R(3,2)*px + R(3,3)*y + yout = R(3,1)*x + R(3,2)*px + R(3,3)*y + R(3,4)*py + R(3,5)*t + R(3,6)*pt; pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y + R(4,4)*py + R(4,5)*t + R(4,6)*pt; - p.pos(RealAoS::t) = R(5,1)*x + R(5,2)*px + R(5,3)*y + tout = R(5,1)*x + R(5,2)*px + R(5,3)*y + R(5,4)*py + R(5,5)*t + R(5,6)*pt; ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y + R(6,4)*py + R(6,5)*t + R(6,6)*pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; diff --git a/src/particles/elements/SoftSol.H b/src/particles/elements/SoftSol.H index bbd7f4d40..f3de4eaf0 100644 --- a/src/particles/elements/SoftSol.H +++ b/src/particles/elements/SoftSol.H @@ -178,7 +178,9 @@ namespace SoftSolenoidData /** This is a soft-edge solenoid functor, so that a variable of this type can be used * like a soft-edge solenoid function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -186,7 +188,9 @@ namespace SoftSolenoidData */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -195,12 +199,10 @@ namespace SoftSolenoidData { using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - - // initialize output values of momenta + // intialize output values + amrex::ParticleReal xout = x; + amrex::ParticleReal yout = y; + amrex::ParticleReal tout = t; amrex::ParticleReal pxout = px; amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; @@ -215,20 +217,23 @@ namespace SoftSolenoidData // so that, e.g., R(3,4) = dyf/dpyi. // push particles using the linear map - p.pos(RealAoS::x) = R(1,1)*x + R(1,2)*px + R(1,3)*y + xout = R(1,1)*x + R(1,2)*px + R(1,3)*y + R(1,4)*py + R(1,5)*t + R(1,6)*pt; pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y + R(2,4)*py + R(2,5)*t + R(2,6)*pt; - p.pos(RealAoS::y) = R(3,1)*x + R(3,2)*px + R(3,3)*y + yout = R(3,1)*x + R(3,2)*px + R(3,3)*y + R(3,4)*py + R(3,5)*t + R(3,6)*pt; pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y + R(4,4)*py + R(4,5)*t + R(4,6)*pt; - p.pos(RealAoS::t) = R(5,1)*x + R(5,2)*px + R(5,3)*y + tout = R(5,1)*x + R(5,2)*px + R(5,3)*y + R(5,4)*py + R(5,5)*t + R(5,6)*pt; ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y + R(6,4)*py + R(6,5)*t + R(6,6)*pt; - // assign updated momenta + // assign updated values + x = xout; + y = yout; + t = tout; px = pxout; py = pyout; pt = ptout; diff --git a/src/particles/elements/Sol.H b/src/particles/elements/Sol.H index 0fb93c6f0..4bee7a09a 100644 --- a/src/particles/elements/Sol.H +++ b/src/particles/elements/Sol.H @@ -49,7 +49,9 @@ namespace impactx /** This is a sol functor, so that a variable of this type can be used like a sol function. * - * @param p Particle AoS data for positions and cpu/id + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t * @param px particle momentum in x * @param py particle momentum in y * @param pt particle momentum in t @@ -57,7 +59,9 @@ namespace impactx */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& AMREX_RESTRICT p, + 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, @@ -65,11 +69,6 @@ namespace impactx using namespace amrex::literals; // for _rt and _prt - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // length of the current slice amrex::ParticleReal const slice_ds = m_ds / nslice(); @@ -79,7 +78,7 @@ namespace impactx // compute phase advance per unit length (in rad/m) and // rotation angle (in rad) - amrex::ParticleReal const alpha = m_ks/2.0_prt; + amrex::ParticleReal const alpha = m_ks / 2.0_prt; amrex::ParticleReal const theta = alpha*slice_ds; // intialize output values @@ -100,26 +99,25 @@ namespace impactx tout = t + (slice_ds/betgam2)*pt; ptout = pt; - // assign intermediate momenta + // assign updated momenta px = pxout; py = pyout; pt = ptout; // advance positions and momenta using map for rotation - p.pos(RealAoS::x) = cos(theta)*xout + sin(theta)*yout; + x = cos(theta)*xout + sin(theta)*yout; pxout = cos(theta)*px + sin(theta)*py; - p.pos(RealAoS::y) = -sin(theta)*xout + cos(theta)*yout; + y = -sin(theta)*xout + cos(theta)*yout; pyout = -sin(theta)*px + cos(theta)*py; - p.pos(RealAoS::t) = tout; + t = tout; ptout = pt; - // assign updated momenta + // assign updated values px = pxout; py = pyout; pt = ptout; - } /** This pushes the reference particle. diff --git a/src/particles/elements/diagnostics/openPMD.cpp b/src/particles/elements/diagnostics/openPMD.cpp index d61695842..212cf1acc 100644 --- a/src/particles/elements/diagnostics/openPMD.cpp +++ b/src/particles/elements/diagnostics/openPMD.cpp @@ -242,19 +242,20 @@ namespace detail // define data set and metadata io::Datatype dtype_fl = io::determineDatatype(); io::Datatype dtype_ui = io::determineDatatype(); + io::Datatype dtype_in = io::determineDatatype(); auto d_fl = io::Dataset(dtype_fl, {np}); auto d_ui = io::Dataset(dtype_ui, {np}); + auto d_in = io::Dataset(dtype_in, {np}); - // AoS: Real + // SoA: Real { - std::vector real_aos_names(RealAoS::names_s.size()); - std::copy(RealAoS::names_s.begin(), RealAoS::names_s.end(), real_aos_names.begin()); - for (auto real_idx=0; real_idx < RealAoS::nattribs; real_idx++) { - auto const component_name = real_aos_names.at(real_idx); + std::vector real_soa_names(RealSoA::names_s.size()); + std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); + for (auto real_idx = 0; real_idx < RealSoA::nattribs; real_idx++) { + auto const component_name = real_soa_names.at(real_idx); getComponentRecord(component_name).resetDataset(d_fl); } } - // openPMD coarse position { beam["positionOffset"]["x"].resetDataset(d_fl); @@ -265,20 +266,20 @@ namespace detail beam["positionOffset"]["t"].makeConstant(ref_part.t); } - // AoS: Int - beam["id"][scalar].resetDataset(d_ui); - - // SoA: Real + // SoA: Int { - std::vector real_soa_names(RealSoA::names_s.size()); - std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); - for (auto real_idx = 0; real_idx < RealSoA::nattribs; real_idx++) { - auto const component_name = real_soa_names.at(real_idx); - getComponentRecord(component_name).resetDataset(d_fl); + // special case: id & cpu to openPMD int64 id + beam["id"][scalar].resetDataset(d_ui); + + // all other int properties + int const skip = 2; // skip: id, cpu + std::vector int_soa_names(IntSoA::names_s.size()); + std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_soa_names.begin()); + for (auto int_idx = skip; int_idx < IntSoA::nattribs; int_idx++) { + auto const component_name = int_soa_names.at(int_idx); + getComponentRecord(component_name).resetDataset(d_in); } } - // SoA: Int - static_assert(IntSoA::nattribs == 0); // not yet used #else amrex::ignore_unused(pc, step); #endif @@ -348,9 +349,6 @@ namespace detail auto & offset = m_offset.at(currentLevel); // ... - // preparing access to particle data: AoS - auto& aos = pti.GetArrayOfStructs(); - // series & iteration auto series = std::any_cast(m_series); io::WriteIterations iterations = series.writeIterations(); @@ -371,35 +369,7 @@ namespace detail return detail::get_component_record(beam, comp_name); }; - // AoS: position and particle ID - { - using vs = std::vector; - vs const positionComponents{"x", "y", "t"}; // TODO: generalize - for (auto currDim = 0; currDim < AMREX_SPACEDIM; currDim++) { - std::shared_ptr curr( - new amrex::ParticleReal[numParticleOnTile], - [](amrex::ParticleReal const *p) { delete[] p; } - ); - for (auto i = 0; i < numParticleOnTile; i++) { - curr.get()[i] = aos[i].pos(currDim); - } - std::string const positionComponent = positionComponents[currDim]; - beam["position"][positionComponent].storeChunk(curr, {offset}, - {numParticleOnTile64}); - } - - // save particle ID after converting it to a globally unique ID - std::shared_ptr ids( - new uint64_t[numParticleOnTile], - [](uint64_t const *p) { delete[] p; } - ); - for (auto i = 0; i < numParticleOnTile; i++) { - ids.get()[i] = ablastr::particles::localIDtoGlobal(aos[i].id(), aos[i].cpu()); - } - beam["id"][scalar].storeChunk(ids, {offset}, {numParticleOnTile64}); - } - - // SoA: everything else + // SoA auto const& soa = pti.GetStructOfArrays(); // SoA floating point (ParticleReal) properties { @@ -412,21 +382,31 @@ namespace detail soa.GetRealData(real_idx).data(), {offset}, {numParticleOnTile64}); } } - // SoA integer (int) properties (not yet used) + // SoA integer (int) properties { - static_assert(IntSoA::nattribs == 0); // not yet used - /* - // comment this in once IntSoA::nattribs is > 0 + // special case: id & cpu to openPMD int64 id + std::shared_ptr ids( + new uint64_t[numParticleOnTile], + [](uint64_t const *p) { delete[] p; } + ); + for (auto i = 0; i < numParticleOnTile; i++) { + ids.get()[i] = ablastr::particles::localIDtoGlobal( + soa.GetIntData(IntSoA::id).data()[i], + soa.GetIntData(IntSoA::cpu).data()[i] + ); + } + beam["id"][scalar].storeChunk(ids, {offset}, {numParticleOnTile64}); - std::vector int_soa_names(IntSoA::names_s.size); + // all other int properties + int const skip = 2; // skip: id, cpu + std::vector int_soa_names(IntSoA::names_s.size()); std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_soa_names.begin()); - for (auto int_idx=0; int_idx < RealSoA::nattribs; int_idx++) { + for (auto int_idx = skip; int_idx < IntSoA::nattribs; int_idx++) { auto const component_name = int_soa_names.at(int_idx); getComponentRecord(component_name).storeChunkRaw( soa.GetIntData(int_idx).data(), {offset}, {numParticleOnTile64}); } - */ } // TODO diff --git a/src/particles/elements/mixin/beamoptic.H b/src/particles/elements/mixin/beamoptic.H index 4461558bd..f48765e7e 100644 --- a/src/particles/elements/mixin/beamoptic.H +++ b/src/particles/elements/mixin/beamoptic.H @@ -40,25 +40,31 @@ namespace detail struct PushSingleParticle { using PType = ImpactXParticleContainer::ParticleType; + using ParticleTileType = amrex::ParticleTile,RealSoA::nattribs, IntSoA::nattribs>; /** Constructor taking in pointers to particle data * * @param element the beamline element to push through - * @param aos_ptr the array-of-struct with position and ids + * @param part_x the array to the particle position (x) + * @param part_y the array to the particle position (y) + * @param part_t the array to the particle position (t) * @param part_px the array to the particle momentum (x) * @param part_py the array to the particle momentum (y) * @param part_pt the array to the particle momentum (t) * @param ref_part the struct containing the reference particle */ PushSingleParticle (T_Element element, - PType* AMREX_RESTRICT aos_ptr, + amrex::ParticleReal* AMREX_RESTRICT part_x, + amrex::ParticleReal* AMREX_RESTRICT part_y, + amrex::ParticleReal* AMREX_RESTRICT part_t, amrex::ParticleReal* AMREX_RESTRICT part_px, amrex::ParticleReal* AMREX_RESTRICT part_py, amrex::ParticleReal* AMREX_RESTRICT part_pt, RefPart ref_part) - : m_element(std::move(element)), m_aos_ptr(aos_ptr), - m_part_px(part_px), m_part_py(part_py), m_part_pt(part_pt), - m_ref_part(std::move(ref_part)) + : m_element(std::move(element)), + m_part_x(part_x), m_part_y(part_y), m_part_t(part_t), + m_part_px(part_px), m_part_py(part_py), m_part_pt(part_pt), + m_ref_part(std::move(ref_part)) { } @@ -75,22 +81,24 @@ namespace detail void operator() (long i) const { - // access AoS data such as positions and cpu/id - PType& AMREX_RESTRICT p = m_aos_ptr[i]; - // access SoA Real data + amrex::ParticleReal & AMREX_RESTRICT x = m_part_x[i]; + amrex::ParticleReal & AMREX_RESTRICT y = m_part_y[i]; + amrex::ParticleReal & AMREX_RESTRICT t = m_part_t[i]; amrex::ParticleReal & AMREX_RESTRICT px = m_part_px[i]; amrex::ParticleReal & AMREX_RESTRICT py = m_part_py[i]; amrex::ParticleReal & AMREX_RESTRICT pt = m_part_pt[i]; // push through element - m_element(p, px, py, pt, m_ref_part); + m_element(x, y, t, px, py, pt, m_ref_part); } private: T_Element const m_element; - PType* const AMREX_RESTRICT m_aos_ptr; + amrex::ParticleReal* const AMREX_RESTRICT m_part_x; + amrex::ParticleReal* const AMREX_RESTRICT m_part_y; + amrex::ParticleReal* const AMREX_RESTRICT m_part_t; amrex::ParticleReal* const AMREX_RESTRICT m_part_px; amrex::ParticleReal* const AMREX_RESTRICT m_part_py; amrex::ParticleReal* const AMREX_RESTRICT m_part_pt; @@ -107,19 +115,17 @@ namespace detail ) { const int np = pti.numParticles(); - // preparing access to particle data: AoS - using PType = ImpactXParticleContainer::ParticleType; - auto& aos = pti.GetArrayOfStructs(); - PType* AMREX_RESTRICT aos_ptr = aos().dataPtr(); - // preparing access to particle data: SoA of Reals auto& soa_real = pti.GetStructOfArrays().GetRealData(); + amrex::ParticleReal* const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_t = soa_real[RealSoA::t].dataPtr(); amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr(); amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr(); amrex::ParticleReal* const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr(); detail::PushSingleParticle const pushSingleParticle( - element, aos_ptr, part_px, part_py, part_pt, ref_part); + element, part_x, part_y, part_t, part_px, part_py, part_pt, ref_part); // loop over beam particles in the box amrex::ParallelFor(np, pushSingleParticle); } diff --git a/src/particles/spacecharge/GatherAndPush.cpp b/src/particles/spacecharge/GatherAndPush.cpp index 4a85abebf..4c451d46f 100644 --- a/src/particles/spacecharge/GatherAndPush.cpp +++ b/src/particles/spacecharge/GatherAndPush.cpp @@ -63,13 +63,11 @@ namespace impactx::spacecharge amrex::ParticleReal const dt = slice_ds / pc.GetRefParticle().beta() / c0_SI; - // preparing access to particle data: AoS - using PType = ImpactXParticleContainer::ParticleType; - auto const & aos = pti.GetArrayOfStructs(); - PType const * const AMREX_RESTRICT aos_ptr = aos().dataPtr(); - // preparing access to particle data: SoA of Reals auto& soa_real = pti.GetStructOfArrays().GetRealData(); + amrex::ParticleReal* const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_z = soa_real[RealSoA::z].dataPtr(); // note: currently for a fixed t amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr(); amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr(); amrex::ParticleReal* const AMREX_RESTRICT part_pz = soa_real[RealSoA::pz].dataPtr(); // note: currently for a fixed t @@ -79,10 +77,10 @@ namespace impactx::spacecharge // gather to each particle and push momentum amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) { - // access AoS data such as positions and cpu/id - PType const & AMREX_RESTRICT p = aos_ptr[i]; - // access SoA Real data + amrex::ParticleReal & AMREX_RESTRICT x = part_x[i]; + amrex::ParticleReal & AMREX_RESTRICT y = part_y[i]; + amrex::ParticleReal & AMREX_RESTRICT z = part_z[i]; amrex::ParticleReal & AMREX_RESTRICT px = part_px[i]; amrex::ParticleReal & AMREX_RESTRICT py = part_py[i]; amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i]; @@ -90,7 +88,7 @@ namespace impactx::spacecharge // force gather amrex::GpuArray const field_interp = ablastr::particles::doGatherVectorFieldNodal ( - p.pos(RealAoS::x), p.pos(RealAoS::y), p.pos(RealAoS::t), + x, y, z, scf_arr_x, scf_arr_y, scf_arr_z, invdr, prob_lo); diff --git a/src/particles/transformation/CoordinateTransformation.cpp b/src/particles/transformation/CoordinateTransformation.cpp index 7615354ed..61fce0879 100644 --- a/src/particles/transformation/CoordinateTransformation.cpp +++ b/src/particles/transformation/CoordinateTransformation.cpp @@ -43,19 +43,17 @@ namespace transformation { for (ParIt pti(pc, lev); pti.isValid(); ++pti) { const int np = pti.numParticles(); - // preparing access to particle data: AoS - using PType = ImpactXParticleContainer::ParticleType; - auto &aos = pti.GetArrayOfStructs(); - PType *AMREX_RESTRICT aos_ptr = aos().dataPtr(); - // preparing access to particle data: SoA of Reals auto &soa_real = pti.GetStructOfArrays().GetRealData(); + amrex::ParticleReal *const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr(); + amrex::ParticleReal *const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr(); amrex::ParticleReal *const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr(); amrex::ParticleReal *const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr(); if( direction == Direction::to_fixed_s) { BL_PROFILE("impactx::transformation::CoordinateTransformation::to_fixed_s"); + amrex::ParticleReal *const AMREX_RESTRICT part_z = soa_real[RealSoA::z].dataPtr(); amrex::ParticleReal *const AMREX_RESTRICT part_pz = soa_real[RealSoA::pz].dataPtr(); // Design value of pz/mc = beta*gamma @@ -63,33 +61,34 @@ namespace transformation { ToFixedS const to_s(pzd); amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(long i) { - // access AoS data such as positions and cpu/id - PType &p = aos_ptr[i]; - - // access SoA Real data + amrex::ParticleReal &x = part_x[i]; + amrex::ParticleReal &y = part_y[i]; + amrex::ParticleReal &z = part_z[i]; amrex::ParticleReal &px = part_px[i]; amrex::ParticleReal &py = part_py[i]; amrex::ParticleReal &pz = part_pz[i]; - to_s(p, px, py, pz); + to_s(x, y, z, px, py, pz); }); } else { BL_PROFILE("impactx::transformation::CoordinateTransformation::to_fixed_t"); + amrex::ParticleReal *const AMREX_RESTRICT part_t = soa_real[RealSoA::t].dataPtr(); amrex::ParticleReal *const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr(); - amrex::ParticleReal const ptd = pd; // Design value of pt/mc2 = -gamma. + // Design value of pt/mc2 = -gamma. + amrex::ParticleReal const ptd = pd; + ToFixedT const to_t(ptd); amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(long i) { - // access AoS data such as positions and cpu/id - PType &p = aos_ptr[i]; - - // access SoA Real data + amrex::ParticleReal &x = part_x[i]; + amrex::ParticleReal &y = part_y[i]; + amrex::ParticleReal &t = part_t[i]; amrex::ParticleReal &px = part_px[i]; amrex::ParticleReal &py = part_py[i]; amrex::ParticleReal &pt = part_pt[i]; - to_t(p, px, py, pt); + to_t(x, y, t, px, py, pt); }); } } // end loop over all particle boxes diff --git a/src/particles/transformation/ToFixedS.H b/src/particles/transformation/ToFixedS.H index 859f08807..1ce00566d 100644 --- a/src/particles/transformation/ToFixedS.H +++ b/src/particles/transformation/ToFixedS.H @@ -42,25 +42,24 @@ namespace transformation /** This is a t-to-s map, so that a variable of this type can be used like a * t-to-s function. * - * @param[inout] p Particle AoS data for positions and cpu/id + * @param[inout] x particle position in x + * @param[inout] y particle position in y + * @param[inout] z particle position in z (in), in t (out) * @param[inout] px particle momentum in x * @param[inout] py particle momentum in y * @param[inout] pz particle momentum in z (in), in t (out) */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& p, + amrex::ParticleReal & x, + amrex::ParticleReal & y, + amrex::ParticleReal & z, amrex::ParticleReal & px, amrex::ParticleReal & py, amrex::ParticleReal & pz) const { using namespace amrex::literals; - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const z = p.pos(RealAoS::z); - // compute value of reference ptd = -gamma amrex::ParticleReal const argd = 1.0_prt + pow(m_pzd, 2); AMREX_ASSERT_WITH_MESSAGE(argd > 0.0_prt, "invalid ptd arg (<=0)"); @@ -78,11 +77,12 @@ namespace transformation amrex::ParticleReal const ptf = arg > 0.0_prt ? -sqrt(arg) : -1.0_prt; // transform position and momentum (from fixed t to fixed s) - p.pos(RealAoS::x) = x - px * z / (m_pzd + pz); + x = x - px * z / (m_pzd + pz); // px = px; - p.pos(RealAoS::y) = y - py * z / (m_pzd + pz); + y = y - py * z / (m_pzd + pz); // py = py; - p.pos(RealAoS::t) = ptf * z / (m_pzd + pz); + auto & t = z; // We store t in the same memory slot as z. + t = ptf * z / (m_pzd + pz); auto & pt = pz; // We store pt in the same memory slot as pz. pt = ptf - ptdf; diff --git a/src/particles/transformation/ToFixedT.H b/src/particles/transformation/ToFixedT.H index 1515075b5..85c26d65f 100644 --- a/src/particles/transformation/ToFixedT.H +++ b/src/particles/transformation/ToFixedT.H @@ -42,25 +42,24 @@ namespace transformation /** This is a s-to-t map, so that a variable of this type can be used like a * s-to-t function. * - * @param[inout] p Particle AoS data for positions and cpu/id + * @param[inout] x particle position in x + * @param[inout] y particle position in y + * @param[inout] t particle position in t (in), in z (out) * @param[inout] px particle momentum in x * @param[inout] py particle momentum in y * @param[inout] pt particle momentum in t (in), in z (out) */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( - PType& p, + amrex::ParticleReal & x, + amrex::ParticleReal & y, + amrex::ParticleReal & t, amrex::ParticleReal & px, amrex::ParticleReal & py, amrex::ParticleReal & pt) const { using namespace amrex::literals; - // access AoS data such as positions and cpu/id - amrex::ParticleReal const x = p.pos(RealAoS::x); - amrex::ParticleReal const y = p.pos(RealAoS::y); - amrex::ParticleReal const t = p.pos(RealAoS::t); - // compute value of reference pzd = beta*gamma amrex::ParticleReal const argd = -1.0_prt + pow(m_ptd, 2); AMREX_ASSERT_WITH_MESSAGE(argd > 0.0_prt, "invalid pzd arg (<=0)"); @@ -78,11 +77,12 @@ namespace transformation amrex::ParticleReal const pzf = arg > 0.0_prt ? sqrt(arg) : 0.0_prt; // transform position and momentum (from fixed s to fixed t) - p.pos(RealAoS::x) = x + px*t/(m_ptd+pt); + x = x + px*t/(m_ptd+pt); // px = px; - p.pos(RealAoS::y) = y + py*t/(m_ptd+pt); + y = y + py*t/(m_ptd+pt); // py = py; - p.pos(RealAoS::z) = pzf * t / (m_ptd + pt); + auto & z = t; // We store z in the same memory slot as t. + z = pzf * t / (m_ptd + pt); auto & pz = pt; // We store pz in the same memory slot as pt. pz = pzf - pzdf; diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp index dae8d2c71..7217471a2 100644 --- a/src/python/ImpactXParticleContainer.cpp +++ b/src/python/ImpactXParticleContainer.cpp @@ -17,28 +17,28 @@ using namespace impactx; void init_impactxparticlecontainer(py::module& m) { py::class_< - ParIter, - amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> + ParIterSoA, + amrex::ParIterSoA >(m, "ImpactXParIter") - .def(py::init(), + .def(py::init(), py::arg("particle_container"), py::arg("level")) - .def(py::init(), + .def(py::init(), py::arg("particle_container"), py::arg("level"), py::arg("info")) ; py::class_< - ParConstIter, - amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> + ParConstIterSoA, + amrex::ParConstIterSoA >(m, "ImpactXParConstIter") - .def(py::init(), + .def(py::init(), py::arg("particle_container"), py::arg("level")) - .def(py::init(), + .def(py::init(), py::arg("particle_container"), py::arg("level"), py::arg("info")) ; py::class_< ImpactXParticleContainer, - amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs> + amrex::ParticleContainerPureSoA >(m, "ImpactXParticleContainer") //.def(py::init<>()) .def("add_n_particles",