From 21b43c38ec2e6907a9aa1ce7adf15524b7f35f46 Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Tue, 13 Feb 2024 22:21:15 +0100 Subject: [PATCH 01/11] use id is_valid --- src/particles/beam/BeamParticleContainer.cpp | 5 ++--- src/particles/beam/BeamParticleContainerInit.cpp | 1 + src/particles/deposition/BeamDepositCurrent.cpp | 2 +- src/particles/deposition/ExplicitDeposition.cpp | 2 +- src/particles/deposition/PlasmaDepositCurrent.cpp | 4 ++-- src/particles/plasma/PlasmaParticleContainer.cpp | 2 +- src/particles/pusher/BeamParticleAdvance.cpp | 2 +- src/particles/pusher/GetAndSetPosition.H | 4 ++-- src/particles/pusher/PlasmaParticleAdvance.cpp | 2 +- src/salame/Salame.cpp | 2 +- 10 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/particles/beam/BeamParticleContainer.cpp b/src/particles/beam/BeamParticleContainer.cpp index 9df96bf49c..4fd091f60f 100644 --- a/src/particles/beam/BeamParticleContainer.cpp +++ b/src/particles/beam/BeamParticleContainer.cpp @@ -378,8 +378,7 @@ BeamParticleContainer::intializeSlice (int slice, int which_slice) { ptd.rdata(BeamIdx::uy)[ip] = ptd_init.rdata(BeamIdx::uy)[idx_src]; ptd.rdata(BeamIdx::uz)[ip] = ptd_init.rdata(BeamIdx::uz)[idx_src]; - ptd.id(ip) = ptd_init.id(idx_src); - ptd.cpu(ip) = 0; + ptd.idcpu(ip) = ptd_init.idcpu(idx_src); ptd.idata(BeamIdx::nsubcycles)[ip] = 0; } ); @@ -472,7 +471,7 @@ BeamParticleContainer::InSituComputeDiags (int islice) const amrex::Real uz_inv = uz == 0._rt ? 0._rt : 1._rt / uz; - if (ptd.id(ip) < 0 || x*x + y*y > insitu_radius_sq) { + if (!ptd.id(ip).is_valid() || x*x + y*y > insitu_radius_sq) { return amrex::IdentityTuple(ReduceTuple{}, reduce_op); } const amrex::Real gamma = std::sqrt(1.0_rt + ux*ux + uy*uy + uz*uz); diff --git a/src/particles/beam/BeamParticleContainerInit.cpp b/src/particles/beam/BeamParticleContainerInit.cpp index 20af4ea94d..4b043a3d89 100644 --- a/src/particles/beam/BeamParticleContainerInit.cpp +++ b/src/particles/beam/BeamParticleContainerInit.cpp @@ -53,6 +53,7 @@ namespace ptd.rdata(BeamIdx::w )[ip] = std::abs(weight); ptd.id(ip) = pid > 0 ? pid + ip : pid; + ptd.cpu(ip) = 0; } /** \brief Adds a single beam particle into the per-slice BeamTile diff --git a/src/particles/deposition/BeamDepositCurrent.cpp b/src/particles/deposition/BeamDepositCurrent.cpp index 41f13733a1..494d2803ae 100644 --- a/src/particles/deposition/BeamDepositCurrent.cpp +++ b/src/particles/deposition/BeamDepositCurrent.cpp @@ -99,7 +99,7 @@ DepositCurrentSlice (BeamParticleContainer& beam, Fields& fields, // Skip invalid particles and ghost particles not in the last slice // beam deposits only up to its finest level - if (ptd.id(ip) < 0 || (only_highest ? (ptd.cpu(ip)!=lev) : (ptd.cpu(ip) insitu_radius_sq) { + if (!ptd.id(ip).is_valid() || x*x + y*y > insitu_radius_sq) { return amrex::IdentityTuple(ReduceTuple{}, reduce_op); } // particle's Lorentz factor diff --git a/src/particles/pusher/BeamParticleAdvance.cpp b/src/particles/pusher/BeamParticleAdvance.cpp index 46de0d9c9e..4481ebd58a 100644 --- a/src/particles/pusher/BeamParticleAdvance.cpp +++ b/src/particles/pusher/BeamParticleAdvance.cpp @@ -129,7 +129,7 @@ AdvanceBeamParticlesSlice ( beam.getNumParticlesIncludingSlipped(WhichBeamSlice::This), [=] AMREX_GPU_DEVICE (int ip, auto depos_order, auto c_use_external_fields) { - if (ptd.id(ip) < 0) return; + if (!ptd.id(ip).is_valid()) return; amrex::Real xp = ptd.pos(0, ip); amrex::Real yp = ptd.pos(1, ip); diff --git a/src/particles/pusher/GetAndSetPosition.H b/src/particles/pusher/GetAndSetPosition.H index 5c065a0221..795041713f 100644 --- a/src/particles/pusher/GetAndSetPosition.H +++ b/src/particles/pusher/GetAndSetPosition.H @@ -72,7 +72,7 @@ struct EnforceBCandSetPos const bool invalid = (shifted && !m_is_per); if (invalid) { ptd.rdata(PlasmaIdx::w)[ip] = 0.0_rt; - p.id() = -std::abs(p.id()); + p.id().make_invalid(); } x = p.pos(0); y = p.pos(1); @@ -103,7 +103,7 @@ struct EnforceBCandSetPos const bool invalid = (shifted && !m_is_per); if (invalid) { ptd.rdata(BeamIdx::w)[ip] = 0.0_rt; - p.id() = -std::abs(p.id()); + p.id().make_invalid(); } x = p.pos(0); y = p.pos(1); diff --git a/src/particles/pusher/PlasmaParticleAdvance.cpp b/src/particles/pusher/PlasmaParticleAdvance.cpp index a3b7501ec0..c72b45bb33 100644 --- a/src/particles/pusher/PlasmaParticleAdvance.cpp +++ b/src/particles/pusher/PlasmaParticleAdvance.cpp @@ -107,7 +107,7 @@ AdvancePlasmaParticles (PlasmaParticleContainer& plasma, const Fields & fields, [=] AMREX_GPU_DEVICE (int idx, auto depos_order, auto use_laser) { const int ip = idx + idx_begin; // only push plasma particles on their according MR level - if (ptd.id(ip) < 0 || ptd.cpu(ip) != lev) return; + if (!ptd.id(ip).is_valid() || ptd.cpu(ip) != lev) return; // define field at particle position reals amrex::Real ExmByp = 0._rt, EypBxp = 0._rt, Ezp = 0._rt; diff --git a/src/salame/Salame.cpp b/src/salame/Salame.cpp index 79223f33b1..894ad8338d 100644 --- a/src/salame/Salame.cpp +++ b/src/salame/Salame.cpp @@ -302,7 +302,7 @@ SalameOnlyAdvancePlasma (Hipace* hipace, const int lev) [=] AMREX_GPU_DEVICE (long idx, auto depos_order) { const int ip = idx + idx_begin; // only push plasma particles on their according MR level - if (ptd.id(ip) < 0 || ptd.cpu(ip) != lev) return; + if (!ptd.id(ip).is_valid() || ptd.cpu(ip) != lev) return; const amrex::Real xp = ptd.rdata(PlasmaIdx::x_prev)[ip]; const amrex::Real yp = ptd.rdata(PlasmaIdx::y_prev)[ip]; From 3f04bad3c89964894c792022912d1f08aa53d76e Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Tue, 13 Feb 2024 23:03:32 +0100 Subject: [PATCH 02/11] use amrex removeInvalidParticles for the beam --- src/particles/sorting/SliceSort.cpp | 93 ++++++----------------------- 1 file changed, 17 insertions(+), 76 deletions(-) diff --git a/src/particles/sorting/SliceSort.cpp b/src/particles/sorting/SliceSort.cpp index 02ebb0dc76..59d9010308 100644 --- a/src/particles/sorting/SliceSort.cpp +++ b/src/particles/sorting/SliceSort.cpp @@ -19,39 +19,21 @@ shiftSlippedParticles (BeamParticleContainer& beam, const int slice, amrex::Geom return; } - const auto ptd = beam.getBeamSlice(WhichBeamSlice::This).getParticleTileData(); + // remove all invalid particles from WhichBeamSlice::This (including slipped) + amrex::removeInvalidParticles(beam.getBeamSlice(WhichBeamSlice::This)); // min_z is the lower end of WhichBeamSlice::This const amrex::Real min_z = geom.ProbLo(2) + (slice-geom.Domain().smallEnd(2))*geom.CellSize(2); - amrex::ReduceOps reduce_op; - amrex::ReduceData reduce_data(reduce_op); - using ReduceTuple = typename decltype(reduce_data)::Type; - - const int num_particles = beam.getNumParticlesIncludingSlipped(WhichBeamSlice::This); - - // count the number of invalid and slipped particles - reduce_op.eval( - num_particles, reduce_data, - [=] AMREX_GPU_DEVICE (const int ip) -> ReduceTuple - { - if (ptd.id(ip) < 0) { - return {1, 0}; - } else if (ptd.pos(2, ip) < min_z) { - return {0, 1}; - } else { - return {0, 0}; - } + // put non slipped particles at the start of the slice + const int num_stay = amrex::partitionParticles(beam.getBeamSlice(WhichBeamSlice::This), + [=] AMREX_GPU_DEVICE (auto& ptd, int i) { + return ptd.pos(2, i) >= min_z; }); - ReduceTuple t = reduce_data.value(); - - const int num_invalid = amrex::get<0>(t); - const int num_slipped = amrex::get<1>(t); - const int num_stay = beam.getNumParticlesIncludingSlipped(WhichBeamSlice::This) - - num_invalid - num_slipped; + const int num_slipped = beam.getBeamSlice(WhichBeamSlice::This).size() - num_stay; - if (num_invalid == 0 && num_slipped == 0) { + if (num_slipped == 0) { // nothing to do beam.resize(WhichBeamSlice::This, num_stay, 0); return; @@ -64,60 +46,19 @@ shiftSlippedParticles (BeamParticleContainer& beam, const int slice, amrex::Geom beam.resize(WhichBeamSlice::Next, next_size, num_slipped); - BeamTile tmp{}; - tmp.resize(num_stay); - - const auto ptd_tmp = tmp.getParticleTileData(); - - // copy valid non slipped particles to the tmp tile - const int num_stay2 = amrex::Scan::PrefixSum (num_particles, - [=] AMREX_GPU_DEVICE (const int ip) -> int - { - return ptd.id(ip) >= 0 && ptd.pos(2, ip) >= min_z; - }, - [=] AMREX_GPU_DEVICE (const int ip, const int s) - { - if (ptd.id(ip) >= 0 && ptd.pos(2, ip) >= min_z) { - ptd_tmp.idcpu(s) = ptd.idcpu(ip); - for (int j=0; j (num_particles, - [=] AMREX_GPU_DEVICE (const int ip) -> int + amrex::ParallelFor(num_slipped, + [=] AMREX_GPU_DEVICE (int i) { - return ptd.id(ip) >= 0 && ptd.pos(2, ip) < min_z; - }, - [=] AMREX_GPU_DEVICE (const int ip, const int s) - { - if (ptd.id(ip) >= 0 && ptd.pos(2, ip) < min_z) { - ptd_next.idcpu(s+next_size) = ptd.idcpu(ip); - for (int j=0; j Date: Wed, 14 Feb 2024 05:37:15 +0100 Subject: [PATCH 03/11] prepare beam communication for runtime components --- src/particles/beam/BeamParticleContainer.H | 7 + src/particles/beam/BeamParticleContainer.cpp | 17 +- src/utils/MultiBuffer.H | 17 +- src/utils/MultiBuffer.cpp | 215 ++++++++++--------- 4 files changed, 143 insertions(+), 113 deletions(-) diff --git a/src/particles/beam/BeamParticleContainer.H b/src/particles/beam/BeamParticleContainer.H index e97f7d420f..93257d2379 100644 --- a/src/particles/beam/BeamParticleContainer.H +++ b/src/particles/beam/BeamParticleContainer.H @@ -191,6 +191,13 @@ public: return m_total_num_particles; } + int numRealComponents () const { return BeamIdx::real_nattribs; } + int numIntComponents () const { return BeamIdx::int_nattribs; } + + bool communicateIdCpuComponent () const { return true; } + bool communicateRealComponent (int rcomp) const { return rcomp < BeamIdx::real_nattribs_in_buffer; } + bool communicateIntComponent (int icomp) const { return icomp < BeamIdx::int_nattribs_in_buffer; } + private: int m_slice_permutation = 0; std::array m_slices {}; diff --git a/src/particles/beam/BeamParticleContainer.cpp b/src/particles/beam/BeamParticleContainer.cpp index 9df96bf49c..84c88a29ef 100644 --- a/src/particles/beam/BeamParticleContainer.cpp +++ b/src/particles/beam/BeamParticleContainer.cpp @@ -95,11 +95,11 @@ BeamParticleContainer::ReadParameters () auto& soa = getBeamInitSlice().GetStructOfArrays(); soa.GetIdCPUData().setArena( m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena()); - for (int rcomp = 0; rcomp < BeamIdx::real_nattribs_in_buffer; ++rcomp) { + for (int rcomp = 0; rcomp < soa.NumRealComps(); ++rcomp) { soa.GetRealData()[rcomp].setArena( m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena()); } - for (int icomp = 0; icomp < BeamIdx::int_nattribs_in_buffer; ++icomp) { + for (int icomp = 0; icomp < soa.NumIntComps(); ++icomp) { soa.GetIntData()[icomp].setArena( m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena()); } @@ -409,11 +409,12 @@ BeamParticleContainer::ReorderParticles (int beam_slice, int step, amrex::Geomet amrex::PermutationForDeposition(perm, np, ptile, slice_geom.Domain(), slice_geom, m_reorder_idx_type); const unsigned int* permutations = perm.dataPtr(); + auto& soa = ptile.GetStructOfArrays(); { // Create a scope for the temporary vector below BeamTile::RealVector tmp_real(np_total); - for (int comp = 0; comp < BeamIdx::real_nattribs; ++comp) { - auto src = ptile.GetStructOfArrays().GetRealData(comp).data(); + for (int comp = 0; comp < soa.NumRealComps(); ++comp) { + auto src = soa.GetRealData(comp).data(); amrex::ParticleReal* dst = tmp_real.data(); amrex::ParallelFor(np_total, [=] AMREX_GPU_DEVICE (int i) { @@ -421,13 +422,13 @@ BeamParticleContainer::ReorderParticles (int beam_slice, int step, amrex::Geomet }); amrex::Gpu::streamSynchronize(); - ptile.GetStructOfArrays().GetRealData(comp).swap(tmp_real); + soa.GetRealData(comp).swap(tmp_real); } } BeamTile::IntVector tmp_int(np_total); - for (int comp = 0; comp < BeamIdx::int_nattribs; ++comp) { - auto src = ptile.GetStructOfArrays().GetIntData(comp).data(); + for (int comp = 0; comp < soa.NumIntComps(); ++comp) { + auto src = soa.GetIntData(comp).data(); int* dst = tmp_int.data(); amrex::ParallelFor(np_total, [=] AMREX_GPU_DEVICE (int i) { @@ -435,7 +436,7 @@ BeamParticleContainer::ReorderParticles (int beam_slice, int step, amrex::Geomet }); amrex::Gpu::streamSynchronize(); - ptile.GetStructOfArrays().GetIntData(comp).swap(tmp_int); + soa.GetIntData(comp).swap(tmp_int); } } } diff --git a/src/utils/MultiBuffer.H b/src/utils/MultiBuffer.H index f47f07014e..94c10a85b4 100644 --- a/src/utils/MultiBuffer.H +++ b/src/utils/MultiBuffer.H @@ -62,6 +62,15 @@ private: nprogress, }; + // to specify the input for get_buffer_offset + enum struct offset_type { + beam_idcpu, + beam_real, + beam_int, + laser, + total + }; + // struct to store per-slice data struct DataNode { char* m_buffer = nullptr; @@ -130,11 +139,9 @@ private: // write MultiBeam sizes into the metadata array void write_metadata (int slice, MultiBeam& beams, int beam_slice); - // helper functions to get location of individual arrays inside a buffer - std::size_t get_buffer_offset_idcpu (int slice, int ibeam); - std::size_t get_buffer_offset_real (int slice, int ibeam, int rcomp); - std::size_t get_buffer_offset_int (int slice, int ibeam, int icomp); - std::size_t get_buffer_offset_laser (int slice, int icomp); + // helper function to get location of individual arrays inside a buffer + std::size_t get_buffer_offset (int slice, offset_type type, MultiBeam& beams, + int ibeam, int comp); // copy gpu array into buffer at buffer_offset, either dtoh or dtod void memcpy_to_buffer (int slice, std::size_t buffer_offset, diff --git a/src/utils/MultiBuffer.cpp b/src/utils/MultiBuffer.cpp index 3296670386..a26262a3b6 100644 --- a/src/utils/MultiBuffer.cpp +++ b/src/utils/MultiBuffer.cpp @@ -487,24 +487,11 @@ void MultiBuffer::put_time (amrex::Real time) { } void MultiBuffer::write_metadata (int slice, MultiBeam& beams, int beam_slice) { - std::size_t offset = 0; for (int b = 0; b < m_nbeams; ++b) { - const int num_particles = beams.getBeam(b).getNumParticles(beam_slice); // write number of beam particles (per beam) - get_metadata_location(slice)[b + 1] = num_particles; - offset += ((num_particles + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * sizeof(std::uint64_t); - offset += ((num_particles + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * BeamIdx::real_nattribs_in_buffer * sizeof(amrex::Real); - offset += ((num_particles + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * BeamIdx::int_nattribs_in_buffer * sizeof(int); - } - if (m_use_laser) { - offset += (m_laser_slice_box.numPts() * m_laser_ncomp) * sizeof(amrex::Real); + get_metadata_location(slice)[b + 1] = beams.getBeam(b).getNumParticles(beam_slice); } + std::size_t offset = get_buffer_offset(slice, offset_type::total, beams, 0, 0); // write total buffer size get_metadata_location(slice)[0] = (offset+sizeof(storage_type)-1) / sizeof(storage_type); m_datanodes[slice].m_buffer_size = get_metadata_location(slice)[0]; @@ -513,52 +500,56 @@ void MultiBuffer::write_metadata (int slice, MultiBeam& beams, int beam_slice) { AMREX_ALWAYS_ASSERT(get_metadata_location(slice)[0] < std::numeric_limits::max()); } -std::size_t MultiBuffer::get_buffer_offset_idcpu (int slice, int ibeam) { +std::size_t MultiBuffer::get_buffer_offset (int slice, offset_type type, MultiBeam& beams, + int ibeam, int comp) { std::size_t offset = 0; - for (int b = 0; b < ibeam; ++b) { - offset += ((get_metadata_location(slice)[b + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * sizeof(std::uint64_t); - offset += ((get_metadata_location(slice)[b + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * BeamIdx::real_nattribs_in_buffer * sizeof(amrex::Real); - offset += ((get_metadata_location(slice)[b + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * BeamIdx::int_nattribs_in_buffer * sizeof(int); - } - return offset; -} -std::size_t MultiBuffer::get_buffer_offset_real (int slice, int ibeam, int rcomp) { - std::size_t offset = get_buffer_offset_idcpu(slice, ibeam); - offset += ((get_metadata_location(slice)[ibeam + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * sizeof(std::uint64_t); - offset += ((get_metadata_location(slice)[ibeam + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * rcomp * sizeof(amrex::Real); - return offset; -} + for (int b = 0; b < m_nbeams; ++b) { + auto& beam = beams.getBeam(b); + const int num_particles_rund_up = (get_metadata_location(slice)[b + 1] + + buffer_size_roundup - 1) / buffer_size_roundup * buffer_size_roundup; -std::size_t MultiBuffer::get_buffer_offset_int (int slice, int ibeam, int icomp) { - std::size_t offset = get_buffer_offset_idcpu(slice, ibeam); - offset += ((get_metadata_location(slice)[ibeam + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * sizeof(std::uint64_t); - offset += ((get_metadata_location(slice)[ibeam + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * BeamIdx::real_nattribs_in_buffer * sizeof(amrex::Real); - offset += ((get_metadata_location(slice)[ibeam + 1] + buffer_size_roundup - 1) - / buffer_size_roundup * buffer_size_roundup) - * icomp * sizeof(int); - return offset; -} + if (beam.communicateIdCpuComponent()) { + if (type == offset_type::beam_idcpu && ibeam == b) { + return offset; + } + offset += num_particles_rund_up * sizeof(std::uint64_t); + } + + for (int rcomp = 0; rcomp < beam.numRealComponents(); ++rcomp) { + if (beam.communicateRealComponent(rcomp)) { + if (type == offset_type::beam_real && ibeam == b && rcomp == comp) { + return offset; + } + offset += num_particles_rund_up * sizeof(amrex::Real); + } + } + + for (int icomp = 0; icomp < beam.numIntComponents(); ++icomp) { + if (beam.communicateIntComponent(icomp)) { + if (type == offset_type::beam_int && ibeam == b && icomp == comp) { + return offset; + } + offset += num_particles_rund_up * sizeof(int); + } + } + } + + if (m_use_laser) { + for (int lcomp = 0; lcomp < m_laser_ncomp; ++lcomp) { + if (type == offset_type::laser && lcomp == comp) { + return offset; + } + offset += m_laser_slice_box.numPts() * sizeof(amrex::Real); + } + } + + if (type == offset_type::total) { + return offset; + } -std::size_t MultiBuffer::get_buffer_offset_laser (int slice, int icomp) { - AMREX_ALWAYS_ASSERT(m_use_laser); - std::size_t offset = get_buffer_offset_idcpu(slice, m_nbeams); - offset += (m_laser_slice_box.numPts() * icomp) * sizeof(amrex::Real); - return offset; + amrex::Abort("MultiBuffer::get_buffer_offset invalid argument"); + return 0; } void MultiBuffer::memcpy_to_buffer (int slice, std::size_t buffer_offset, @@ -593,20 +584,30 @@ void MultiBuffer::memcpy_from_buffer (int slice, std::size_t buffer_offset, void MultiBuffer::pack_data (int slice, MultiBeam& beams, MultiLaser& laser, int beam_slice) { for (int b = 0; b < m_nbeams; ++b) { - const int num_particles = beams.getBeam(b).getNumParticles(beam_slice); - auto& soa = beams.getBeam(b).getBeamSlice(beam_slice).GetStructOfArrays(); - memcpy_to_buffer(slice, get_buffer_offset_idcpu(slice, b), - soa.GetIdCPUData().dataPtr(), - num_particles * sizeof(std::uint64_t)); - for (int rcomp = 0; rcomp < BeamIdx::real_nattribs_in_buffer; ++rcomp) { - memcpy_to_buffer(slice, get_buffer_offset_real(slice, b, rcomp), - soa.GetRealData(rcomp).dataPtr(), - num_particles * sizeof(amrex::Real)); - } - for (int icomp = 0; icomp < BeamIdx::int_nattribs_in_buffer; ++icomp) { - memcpy_to_buffer(slice, get_buffer_offset_int(slice, b, icomp), - soa.GetIntData(icomp).dataPtr(), - num_particles * sizeof(int)); + auto& beam = beams.getBeam(b); + const int num_particles = beam.getNumParticles(beam_slice); + auto& soa = beam.getBeamSlice(beam_slice).GetStructOfArrays(); + + if (beam.communicateIdCpuComponent()) { + memcpy_to_buffer(slice, get_buffer_offset(slice, offset_type::beam_idcpu, beams, b, 0), + soa.GetIdCPUData().dataPtr(), + num_particles * sizeof(std::uint64_t)); + } + + for (int rcomp = 0; rcomp < beam.numRealComponents(); ++rcomp) { + if (beam.communicateRealComponent(rcomp)) { + memcpy_to_buffer(slice, get_buffer_offset(slice, offset_type::beam_real, beams, b, rcomp), + soa.GetRealData(rcomp).dataPtr(), + num_particles * sizeof(amrex::Real)); + } + } + + for (int icomp = 0; icomp < beam.numIntComponents(); ++icomp) { + if (beam.communicateIntComponent(icomp)) { + memcpy_to_buffer(slice, get_buffer_offset(slice, offset_type::beam_int, beams, b, icomp), + soa.GetIntData(icomp).dataPtr(), + num_particles * sizeof(int)); + } } } if (m_use_laser) { @@ -614,10 +615,10 @@ void MultiBuffer::pack_data (int slice, MultiBeam& beams, MultiLaser& laser, int const int laser_comp_0_1 = (beam_slice == WhichBeamSlice::Next) ? np1jp2_r : np1j00_r; const int laser_comp_2_3 = (beam_slice == WhichBeamSlice::Next) ? n00jp2_r : n00j00_r; // copy real and imag components in one operation - memcpy_to_buffer(slice, get_buffer_offset_laser(slice, 0), + memcpy_to_buffer(slice, get_buffer_offset(slice, offset_type::laser, beams, 0, 0), laser.getSlices()[0].dataPtr(laser_comp_0_1), 2 * m_laser_slice_box.numPts() * sizeof(amrex::Real)); - memcpy_to_buffer(slice, get_buffer_offset_laser(slice, 2), + memcpy_to_buffer(slice, get_buffer_offset(slice, offset_type::laser, beams, 0, 2), laser.getSlices()[0].dataPtr(laser_comp_2_3), 2 * m_laser_slice_box.numPts() * sizeof(amrex::Real)); } @@ -629,35 +630,49 @@ void MultiBuffer::pack_data (int slice, MultiBeam& beams, MultiLaser& laser, int void MultiBuffer::unpack_data (int slice, MultiBeam& beams, MultiLaser& laser, int beam_slice) { for (int b = 0; b < m_nbeams; ++b) { + auto& beam = beams.getBeam(b); const int num_particles = get_metadata_location(slice)[b + 1]; - beams.getBeam(b).resize(beam_slice, num_particles, 0); - auto& soa = beams.getBeam(b).getBeamSlice(beam_slice).GetStructOfArrays(); - memcpy_from_buffer(slice, get_buffer_offset_idcpu(slice, b), - soa.GetIdCPUData().dataPtr(), - num_particles * sizeof(std::uint64_t)); - for (int rcomp = 0; rcomp < BeamIdx::real_nattribs_in_buffer; ++rcomp) { - memcpy_from_buffer(slice, get_buffer_offset_real(slice, b, rcomp), - soa.GetRealData(rcomp).dataPtr(), - num_particles * sizeof(amrex::Real)); - } - // initialize per-slice-only components to zero - for (int rcomp = BeamIdx::real_nattribs_in_buffer; rcomp Date: Wed, 14 Feb 2024 20:19:18 +0100 Subject: [PATCH 04/11] Add support for spin to openpmd beam output --- src/diagnostics/OpenPMDWriter.H | 5 +++ src/diagnostics/OpenPMDWriter.cpp | 41 +++++++++++++++------- src/particles/beam/BeamParticleContainer.H | 2 ++ src/utils/IOUtil.cpp | 5 +++ 4 files changed, 40 insertions(+), 13 deletions(-) diff --git a/src/diagnostics/OpenPMDWriter.H b/src/diagnostics/OpenPMDWriter.H index 737f9801c2..d00ca508bc 100644 --- a/src/diagnostics/OpenPMDWriter.H +++ b/src/diagnostics/OpenPMDWriter.H @@ -88,6 +88,11 @@ private: "momentum_x", "momentum_y", "momentum_z" }; + /** Named Beam runtime SoA attributes */ + amrex::Vector m_real_names_spin { + "spin_x", "spin_y", "spin_z" + }; + /** vector over levels of openPMD-api Series object for output */ std::unique_ptr< openPMD::Series > m_outputSeries; diff --git a/src/diagnostics/OpenPMDWriter.cpp b/src/diagnostics/OpenPMDWriter.cpp index ab03f59fdd..9562701dde 100644 --- a/src/diagnostics/OpenPMDWriter.cpp +++ b/src/diagnostics/OpenPMDWriter.cpp @@ -17,8 +17,6 @@ OpenPMDWriter::OpenPMDWriter () { - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_names.size() == BeamIdx::real_nattribs, - "List of real names in openPMD Writer class do not match BeamIdx::real_nattribs"); amrex::ParmParse pp("hipace"); queryWithParser(pp, "openpmd_backend", m_openpmd_backend); // pick first available backend if default is chosen @@ -194,7 +192,11 @@ OpenPMDWriter::InitBeamData (MultiBeam& beams, const amrex::Vector< std::string }); } - m_real_beam_data[ibeam].resize(m_real_names.size()); + if (beams.getBeam(ibeam).m_do_spin_tracking) { + m_real_beam_data[ibeam].resize(m_real_names.size() + m_real_names_spin.size()); + } else { + m_real_beam_data[ibeam].resize(m_real_names.size()); + } for (std::size_t idx=0; idx real_names = m_real_names; + if (beam.m_do_spin_tracking) { + real_names.insert(real_names.end(), m_real_names_spin.begin(), m_real_names_spin.end()); + } + // initialize beam IO on first slice AMREX_ALWAYS_ASSERT(m_offset[ibeam] <= beam.getTotalNumParticles()); const uint64_t np_total = m_offset[ibeam]; SetupPos(beam_species, beam, np_total, geom); - SetupRealProperties(beam_species, m_real_names, np_total); + SetupRealProperties(beam_species, real_names, np_total); for (std::size_t idx=0; idx b m_uint64_beam_data[ibeam][idx].get() + m_offset[ibeam]); } + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_beam_data.size() == soa.NumRealComps(), + "List of real names in openPMD Writer class do not match the beam"); + for (std::size_t idx=0; idx(),{np}); - /* we have 4 SoA real attributes: weight, ux, uy, uz */ + /* we have 7 or 10 SoA real attributes: x, y, z, weight, ux, uy, uz, (sx, sy, sz) */ int const NumSoARealAttributes = real_comp_names.size(); std::set< std::string > addedRecords; // add meta-data per record only once @@ -402,11 +410,18 @@ OpenPMDWriter::SetupRealProperties (openPMD::ParticleSpecies& currSpecies, std::tie(std::ignore, newRecord) = addedRecords.insert(record_name); if( newRecord ) { currRecord.setUnitDimension( utils::getUnitDimension(record_name) ); - currRecord.setAttribute( "macroWeighted", 0u ); - if( record_name == "momentum" ) - currRecord.setAttribute( "weightingPower", 1.0 ); - else - currRecord.setAttribute( "weightingPower", 0.0 ); + + if( record_name == "weighting") { + currRecord.setAttribute( "macroWeighted", 1u ); + } else { + currRecord.setAttribute( "macroWeighted", 0u ); + } + + if( record_name == "weighting" || record_name == "momentum" || record_name == "spin") { + currRecord.setAttribute( "weightingPower", 1.0 ); + } else { + currRecord.setAttribute( "weightingPower", 0.0 ); + } } // end if newRecord } // end for NumSoARealAttributes } diff --git a/src/particles/beam/BeamParticleContainer.H b/src/particles/beam/BeamParticleContainer.H index e97f7d420f..2a2b87c5f2 100644 --- a/src/particles/beam/BeamParticleContainer.H +++ b/src/particles/beam/BeamParticleContainer.H @@ -219,6 +219,8 @@ public: amrex::GpuArray, 6> m_external_fields; /** Owns data for m_external_fields */ amrex::Array m_external_fields_parser; + /** If spin tracking is enabled for this beam */ + bool m_do_spin_tracking = false; private: std::string m_name; /**< name of the species */ /** injection type, fixed_width or fixed_ppc */ diff --git a/src/utils/IOUtil.cpp b/src/utils/IOUtil.cpp index 0c3c51ffdb..7154b49c52 100644 --- a/src/utils/IOUtil.cpp +++ b/src/utils/IOUtil.cpp @@ -130,6 +130,11 @@ utils::getUnitDimension ( std::string const & record_name ) {openPMD::UnitDimension::I, -1.}, {openPMD::UnitDimension::T, -2.} }; + else if( record_name == "spin" ) return { + {openPMD::UnitDimension::L, 2.}, + {openPMD::UnitDimension::M, 1.}, + {openPMD::UnitDimension::T, -1.} + }; else return {}; } #endif From 216ca5e61b2a6df5b8536abae272ad612f146710 Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Wed, 14 Feb 2024 20:30:47 +0100 Subject: [PATCH 05/11] fix assert --- src/diagnostics/OpenPMDWriter.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diagnostics/OpenPMDWriter.cpp b/src/diagnostics/OpenPMDWriter.cpp index 9562701dde..1d35e3d95c 100644 --- a/src/diagnostics/OpenPMDWriter.cpp +++ b/src/diagnostics/OpenPMDWriter.cpp @@ -298,7 +298,7 @@ OpenPMDWriter::CopyBeams (MultiBeam& beams, const amrex::Vector< std::string > b m_uint64_beam_data[ibeam][idx].get() + m_offset[ibeam]); } - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_beam_data.size() == soa.NumRealComps(), + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_beam_data[ibeam].size() == soa.NumRealComps(), "List of real names in openPMD Writer class do not match the beam"); for (std::size_t idx=0; idx Date: Wed, 21 Feb 2024 21:00:07 +0100 Subject: [PATCH 06/11] add spin push --- src/particles/beam/BeamParticleContainer.H | 18 ++++++++-- src/particles/beam/BeamParticleContainer.cpp | 19 ++++++++++ src/particles/pusher/BeamParticleAdvance.cpp | 37 ++++++++++++++++++++ 3 files changed, 71 insertions(+), 3 deletions(-) diff --git a/src/particles/beam/BeamParticleContainer.H b/src/particles/beam/BeamParticleContainer.H index 514f6b1e59..5605fd6138 100644 --- a/src/particles/beam/BeamParticleContainer.H +++ b/src/particles/beam/BeamParticleContainer.H @@ -191,12 +191,20 @@ public: return m_total_num_particles; } - int numRealComponents () const { return BeamIdx::real_nattribs; } + int numRealComponents () const { + return BeamIdx::real_nattribs + (m_do_spin_tracking ? 3 : 0); + } int numIntComponents () const { return BeamIdx::int_nattribs; } bool communicateIdCpuComponent () const { return true; } - bool communicateRealComponent (int rcomp) const { return rcomp < BeamIdx::real_nattribs_in_buffer; } - bool communicateIntComponent (int icomp) const { return icomp < BeamIdx::int_nattribs_in_buffer; } + bool communicateRealComponent (int rcomp) const { + // communicate all compile-time and runtime real components + return rcomp < numRealComponents(); + } + bool communicateIntComponent (int icomp) const { + // don't communicate nsubcycles + return icomp < BeamIdx::int_nattribs_in_buffer; + } private: int m_slice_permutation = 0; @@ -228,6 +236,10 @@ public: amrex::Array m_external_fields_parser; /** If spin tracking is enabled for this beam */ bool m_do_spin_tracking = false; + /** Initial spin of all particles */ + amrex::RealVect m_initial_spin = {1, 0, 0,}; + /** The anomalous magnetic moment */ + amrex::Real m_spin_anom = 0.00115965218128; private: std::string m_name; /**< name of the species */ /** injection type, fixed_width or fixed_ppc */ diff --git a/src/particles/beam/BeamParticleContainer.cpp b/src/particles/beam/BeamParticleContainer.cpp index 2703135685..7554447d2c 100644 --- a/src/particles/beam/BeamParticleContainer.cpp +++ b/src/particles/beam/BeamParticleContainer.cpp @@ -103,6 +103,11 @@ BeamParticleContainer::ReadParameters () soa.GetIntData()[icomp].setArena( m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena()); } + queryWithParserAlt(pp, "do_spin_tracking", m_do_spin_tracking, pp_alt); + if (m_do_spin_tracking) { + getWithParserAlt(pp, "initial_spin", m_initial_spin, pp_alt); + queryWithParserAlt(pp, "spin_anom", m_spin_anom, pp_alt); + } } amrex::Real @@ -383,6 +388,20 @@ BeamParticleContainer::intializeSlice (int slice, int which_slice) { } ); } + + if (m_do_spin_tracking) { + auto ptd = getBeamSlice(which_slice).getParticleTileData(); + + const amrex::RealVect initial_spin_norm = m_initial_spin / m_initial_spin.vectorLength(); + + amrex::ParallelFor(getNumParticles(which_slice), + [=] AMREX_GPU_DEVICE (const int ip) { + ptd.m_runtime_rdata[0][ip] = initial_spin_norm[0]; + ptd.m_runtime_rdata[1][ip] = initial_spin_norm[1]; + ptd.m_runtime_rdata[2][ip] = initial_spin_norm[2]; + } + ); + } } void diff --git a/src/particles/pusher/BeamParticleAdvance.cpp b/src/particles/pusher/BeamParticleAdvance.cpp index 4481ebd58a..a834cf63fe 100644 --- a/src/particles/pusher/BeamParticleAdvance.cpp +++ b/src/particles/pusher/BeamParticleAdvance.cpp @@ -32,6 +32,8 @@ AdvanceBeamParticlesSlice ( const amrex::Real dt = Hipace::GetInstance().m_dt / n_subcycles; const amrex::Real background_density_SI = Hipace::m_background_density_SI; const bool normalized_units = Hipace::m_normalized_units; + const bool spin_tracking = beam.m_do_spin_tracking; + const amrex::Real spin_anom = beam.m_spin_anom; if (normalized_units && radiation_reaction) { AMREX_ALWAYS_ASSERT_WITH_MESSAGE(background_density_SI!=0, @@ -140,6 +142,13 @@ AdvanceBeamParticlesSlice ( int i = ptd.idata(BeamIdx::nsubcycles)[ip]; + amrex::RealVect spin {0._rt, 0._rt, 0._rt}; + if (spin_tracking) { + spin[0] = ptd.m_runtime_rdata[0][ip]; + spin[1] = ptd.m_runtime_rdata[1][ip]; + spin[2] = ptd.m_runtime_rdata[2][ip]; + } + for (; i < n_subcycles; i++) { if (zp < min_z) { @@ -214,6 +223,28 @@ AdvanceBeamParticlesSlice ( + uy_intermediate*uy_intermediate + uz_intermediate*uz_intermediate )*inv_c2 ); + if (spin_tracking) { + const amrex::RealVect E {ExmByp + clight*Byp, EypBxp - clight*Bxp, Ezp}; + const amrex::RealVect B {Bxp, Byp, Bzp}; + const amrex::RealVect u {ux_intermediate*inv_clight, uy_intermediate*inv_clight, + uz_intermediate*inv_clight}; + const amrex::RealVect beta = u*gamma_intermediate_inv; + const amrex::Real gamma_inv_p1 = + gamma_intermediate_inv / (1._rt + gamma_intermediate_inv); + + const amrex::RealVect omega = std::abs(charge_mass_ratio) * inv_clight * ( + B * gamma_intermediate_inv - beta.crossProduct(E) * gamma_inv_p1 + + spin_anom * ( + B - gamma_inv_p1 * u * beta.dotProduct(B) - beta.crossProduct(E) + ) + ); + + const amrex::RealVect h = omega * dt * 0.5_rt; + const amrex::RealVect s_prime = spin + h.crossProduct(spin); + const amrex::Real o = 1._rt / (1._rt + h.dotProduct(h)); + spin = o * (s_prime + (h.dotProduct(s_prime) * h + h.crossProduct(s_prime))); + } + amrex::ParticleReal uz_next = uz + dt * charge_mass_ratio * ( Ezp + ( ux_intermediate * Byp - uy_intermediate * Bxp ) * gamma_intermediate_inv ); @@ -300,5 +331,11 @@ AdvanceBeamParticlesSlice ( ptd.rdata(BeamIdx::ux)[ip] = ux; ptd.rdata(BeamIdx::uy)[ip] = uy; ptd.rdata(BeamIdx::uz)[ip] = uz; + + if (spin_tracking) { + ptd.m_runtime_rdata[0][ip] = spin[0]; + ptd.m_runtime_rdata[1][ip] = spin[1]; + ptd.m_runtime_rdata[2][ip] = spin[2]; + } }); } From 718659d6399c652899456c238aa3cca609aec536 Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Thu, 22 Feb 2024 20:14:55 +0100 Subject: [PATCH 07/11] allocate runtime components --- src/particles/beam/BeamParticleContainer.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/particles/beam/BeamParticleContainer.cpp b/src/particles/beam/BeamParticleContainer.cpp index 7554447d2c..6bcb4b9de2 100644 --- a/src/particles/beam/BeamParticleContainer.cpp +++ b/src/particles/beam/BeamParticleContainer.cpp @@ -107,6 +107,10 @@ BeamParticleContainer::ReadParameters () if (m_do_spin_tracking) { getWithParserAlt(pp, "initial_spin", m_initial_spin, pp_alt); queryWithParserAlt(pp, "spin_anom", m_spin_anom, pp_alt); + for (auto& beam_tile : m_slices) { + // Use 3 real and 0 int runtime components + beam_tile.define(3, 0); + } } } From ac3ec961edfb153259b81d1726c21a879f0338ec Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Thu, 22 Feb 2024 20:20:05 +0100 Subject: [PATCH 08/11] fix warning --- src/diagnostics/OpenPMDWriter.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/diagnostics/OpenPMDWriter.cpp b/src/diagnostics/OpenPMDWriter.cpp index 1d35e3d95c..17ba17cf3f 100644 --- a/src/diagnostics/OpenPMDWriter.cpp +++ b/src/diagnostics/OpenPMDWriter.cpp @@ -298,8 +298,9 @@ OpenPMDWriter::CopyBeams (MultiBeam& beams, const amrex::Vector< std::string > b m_uint64_beam_data[ibeam][idx].get() + m_offset[ibeam]); } - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_beam_data[ibeam].size() == soa.NumRealComps(), - "List of real names in openPMD Writer class do not match the beam"); + AMREX_ALWAYS_ASSERT_WITH_MESSAGE( + int(m_real_beam_data[ibeam].size()) == soa.NumRealComps(), + "List of real names in openPMD Writer class does not match the beam"); for (std::size_t idx=0; idx Date: Fri, 1 Mar 2024 12:36:29 +0100 Subject: [PATCH 09/11] add input doc --- docs/source/run/parameters.rst | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/docs/source/run/parameters.rst b/docs/source/run/parameters.rst index aac39be2aa..3d690bb211 100644 --- a/docs/source/run/parameters.rst +++ b/docs/source/run/parameters.rst @@ -1039,3 +1039,19 @@ Whether the energy loss due to classical radiation reaction of beam particles is Whether the beam particles undergo energy loss due to classical radiation reaction. The implemented radiation reaction model is based on this publication: `M. Tamburini et al., NJP 12, 123005 `__ In normalized units, `hipace.background_density_SI` must be specified. + +Spin tracking +------------- + +Track the spin of each beam particle as it is rotated by the electromagnetic fields using the +TBMT model. This will add three extra components to each beam particle to store the spin and output +those as part of the beam diagnostic. + +* `` or beams.do_spin_tracking`` (`bool`) optional (default `0`) + Enable spin tracking + +* `` or beams.initial_spin`` (3 `float`) + Initial spin ``sx sy sz`` of all particles. The length of the three components is normalized to one. + +* `` or beams.spin_anom`` (`bool`) optional (default `0.00115965218128`) + The anomalous magnetic moment. The default value is the moment for electrons. From 31a10e9a871223e91e22153ee7e192abf165f547 Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Thu, 7 Mar 2024 00:06:49 +0100 Subject: [PATCH 10/11] fix units --- src/particles/pusher/BeamParticleAdvance.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/particles/pusher/BeamParticleAdvance.cpp b/src/particles/pusher/BeamParticleAdvance.cpp index a834cf63fe..d6e19970e3 100644 --- a/src/particles/pusher/BeamParticleAdvance.cpp +++ b/src/particles/pusher/BeamParticleAdvance.cpp @@ -232,10 +232,10 @@ AdvanceBeamParticlesSlice ( const amrex::Real gamma_inv_p1 = gamma_intermediate_inv / (1._rt + gamma_intermediate_inv); - const amrex::RealVect omega = std::abs(charge_mass_ratio) * inv_clight * ( - B * gamma_intermediate_inv - beta.crossProduct(E) * gamma_inv_p1 + const amrex::RealVect omega = std::abs(charge_mass_ratio) * ( + B * gamma_intermediate_inv - beta.crossProduct(E) * inv_clight * gamma_inv_p1 + spin_anom * ( - B - gamma_inv_p1 * u * beta.dotProduct(B) - beta.crossProduct(E) + B - gamma_inv_p1 * u * beta.dotProduct(B) - beta.crossProduct(E) * inv_clight ) ); From 0db47c18f89050e45b18adcdceaf9825bc821d0f Mon Sep 17 00:00:00 2001 From: AlexanderSinn Date: Wed, 13 Mar 2024 17:08:49 +0100 Subject: [PATCH 11/11] Add reference to the TBMT model --- docs/source/run/parameters.rst | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/source/run/parameters.rst b/docs/source/run/parameters.rst index 3d690bb211..c07dc1eb63 100644 --- a/docs/source/run/parameters.rst +++ b/docs/source/run/parameters.rst @@ -1044,7 +1044,10 @@ Spin tracking ------------- Track the spin of each beam particle as it is rotated by the electromagnetic fields using the -TBMT model. This will add three extra components to each beam particle to store the spin and output +Thomas-Bargmann-Michel-Telegdi (TBMT) model, see +[Z. Gong et al., Matter and Radiation at Extremes 8.6 (2023), https://doi.org/10.1063/5.0152382] +for the details of the implementation. +This will add three extra components to each beam particle to store the spin and output those as part of the beam diagnostic. * `` or beams.do_spin_tracking`` (`bool`) optional (default `0`)