diff --git a/src/particles/beam/BeamParticleContainer.cpp b/src/particles/beam/BeamParticleContainer.cpp index 82386c5d09..b69efd2c3d 100644 --- a/src/particles/beam/BeamParticleContainer.cpp +++ b/src/particles/beam/BeamParticleContainer.cpp @@ -91,26 +91,29 @@ BeamParticleContainer::ReadParameters () "Tilted beams and correlated energy spreads are only implemented for fixed weight beams"); } queryWithParserAlt(pp, "initialize_on_cpu", m_initialize_on_cpu, pp_alt); + queryWithParserAlt(pp, "do_spin_tracking", m_do_spin_tracking, pp_alt); + if (m_do_spin_tracking) { + if (m_injection_type != "from_file") { + 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); + } + getBeamInitSlice().define(3, 0); + } auto& soa = getBeamInitSlice().GetStructOfArrays(); soa.GetIdCPUData().setArena( m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena()); for (int rcomp = 0; rcomp < soa.NumRealComps(); ++rcomp) { - soa.GetRealData()[rcomp].setArena( + soa.GetRealData(rcomp).setArena( m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena()); } for (int icomp = 0; icomp < soa.NumIntComps(); ++icomp) { - soa.GetIntData()[icomp].setArena( + 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); - for (auto& beam_tile : m_slices) { - // Use 3 real and 0 int runtime components - beam_tile.define(3, 0); - } - } } amrex::Real @@ -368,6 +371,7 @@ BeamParticleContainer::initializeSlice (int slice, int which_slice) { const int slice_offset = m_init_sorter.m_box_offsets_cpu[slice]; const auto permutations = m_init_sorter.m_box_permutations.dataPtr(); + const bool do_spin_tracking = m_do_spin_tracking; amrex::ParallelFor(num_particles, [=] AMREX_GPU_DEVICE (const int ip) { @@ -379,7 +383,11 @@ BeamParticleContainer::initializeSlice (int slice, int which_slice) { ptd.rdata(BeamIdx::ux)[ip] = ptd_init.rdata(BeamIdx::ux)[idx_src]; ptd.rdata(BeamIdx::uy)[ip] = ptd_init.rdata(BeamIdx::uy)[idx_src]; ptd.rdata(BeamIdx::uz)[ip] = ptd_init.rdata(BeamIdx::uz)[idx_src]; - + if (do_spin_tracking) { + ptd.m_runtime_rdata[0][ip] = ptd_init.m_runtime_rdata[0][idx_src]; + ptd.m_runtime_rdata[1][ip] = ptd_init.m_runtime_rdata[1][idx_src]; + ptd.m_runtime_rdata[2][ip] = ptd_init.m_runtime_rdata[2][idx_src]; + } ptd.idcpu(ip) = ptd_init.idcpu(idx_src); ptd.idata(BeamIdx::nsubcycles)[ip] = 0; ptd.idata(BeamIdx::mr_level)[ip] = 0; @@ -387,12 +395,11 @@ BeamParticleContainer::initializeSlice (int slice, int which_slice) { ); } - if (m_do_spin_tracking) { + if (m_do_spin_tracking && m_injection_type != "from_file") { HIPACE_PROFILE("BeamParticleContainer::initializeSpin()"); 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]; diff --git a/src/particles/beam/BeamParticleContainerInit.cpp b/src/particles/beam/BeamParticleContainerInit.cpp index 98d42cd6f2..ffc8224911 100644 --- a/src/particles/beam/BeamParticleContainerInit.cpp +++ b/src/particles/beam/BeamParticleContainerInit.cpp @@ -42,10 +42,12 @@ namespace */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void AddOneBeamParticle ( - const BeamTileInit::ParticleTileDataType& ptd, const amrex::Real& x, - const amrex::Real& y, const amrex::Real& z, const amrex::Real& ux, const amrex::Real& uy, - const amrex::Real& uz, const amrex::Real& weight, const amrex::Long pid, - const amrex::Long ip, const amrex::Real& speed_of_light, const EnforceBC& enforceBC) noexcept + const BeamTileInit::ParticleTileDataType& ptd, + const amrex::Real& x, const amrex::Real& y, const amrex::Real& z, + const amrex::Real& ux, const amrex::Real& uy, const amrex::Real& uz, + const amrex::Real& sx, const amrex::Real& sy, const amrex::Real& sz, + const amrex::Real& weight, const amrex::Long pid, const amrex::Long ip, + const amrex::Real& speed_of_light, const EnforceBC& enforceBC, bool do_spin) noexcept { amrex::Real xp = x; amrex::Real yp = y; @@ -59,6 +61,11 @@ namespace ptd.rdata(BeamIdx::ux )[ip] = uxp; ptd.rdata(BeamIdx::uy )[ip] = uyp; ptd.rdata(BeamIdx::uz )[ip] = uz * speed_of_light; + if (do_spin) { + ptd.m_runtime_rdata[0][ip] = sx; + ptd.m_runtime_rdata[1][ip] = sy; + ptd.m_runtime_rdata[2][ip] = sz; + } ptd.rdata(BeamIdx::w )[ip] = std::abs(weight); ptd.idcpu(ip) = pid + ip; @@ -790,6 +797,7 @@ InitBeamFromFile (const std::string input_file, std::string name_particle =""; std::string name_r ="", name_rx ="", name_ry ="", name_rz =""; std::string name_u ="", name_ux ="", name_uy ="", name_uz =""; + std::string name_s ="", name_sx ="", name_sy ="", name_sz =""; std::string name_m ="", name_mm =""; std::string name_q ="", name_qq =""; std::string name_g ="", name_gg =""; @@ -873,6 +881,24 @@ InitBeamFromFile (const std::string input_file, } } } + else if(units == std::array {2., 1., -1., 0., 0., 0., 0.} && + m_do_spin_tracking) { + // spin + if(physical_quantity.first == "spin") { + name_s = physical_quantity.first; + for( auto const& axes_direction : physical_quantity.second ) { + if(axes_direction.first == "x" || axes_direction.first == "X") { + name_sx = axes_direction.first; + } + if(axes_direction.first == "y" || axes_direction.first == "Y") { + name_sy = axes_direction.first; + } + if(axes_direction.first == "z" || axes_direction.first == "Z") { + name_sz = axes_direction.first; + } + } + } + } } } } @@ -953,15 +979,28 @@ InitBeamFromFile (const std::string input_file, } } + if (m_do_spin_tracking) { + for(std::string name_s_c : {name_sx, name_sy, name_sz}) { + if(!series.iterations[num_iteration].particles[name_particle][name_s].contains(name_s_c)) { + amrex::Abort("Beam input file does not contain " + name_s_c + " coordinate in " + + name_s + " (spin). An attempt to read these was done because " + + "do_spin_tracking is on for at least one beam.\n"); + } + } + } // print the names of the arrays in the file that are used for the simulation if(Hipace::m_verbose >= 3){ amrex::Print() << "Beam Input File '" << input_file << "' in Iteration '" << num_iteration << "' and Paricle '" << name_particle - << "' imported with:\nPositon '" << name_r << "' (coordinates '" << name_rx << "', '" + << "' imported with:\nPosition '" << name_r << "' (coordinates '" << name_rx << "', '" << name_ry << "', '" << name_rz << "')\n" << momentum_type << " '" << name_u << "' (coordinates '" << name_ux << "', '" << name_uy << "', '" << name_uz << "')\n" << weighting_type << " '" << name_w << "' (in '" << name_ww << "')\n"; + if (m_do_spin_tracking) { + amrex::Print() << name_u << "' (coordinates '" << name_ux + << "', '" << name_uy << "', '" << name_uz << "')\n"; + } } auto electrons = series.iterations[num_iteration].particles[name_particle]; @@ -987,6 +1026,15 @@ InitBeamFromFile (const std::string input_file, amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del}; std::shared_ptr u_z_data{ reinterpret_cast( amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del}; + std::shared_ptr s_x_data{ reinterpret_cast( + amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ? + sizeof(input_type)*num_to_add : 0) ), del}; + std::shared_ptr s_y_data{ reinterpret_cast( + amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ? + sizeof(input_type)*num_to_add : 0) ), del}; + std::shared_ptr s_z_data{ reinterpret_cast( + amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ? + sizeof(input_type)*num_to_add: 0 ) ), del}; std::shared_ptr w_w_data{ reinterpret_cast( amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del}; @@ -997,6 +1045,11 @@ InitBeamFromFile (const std::string input_file, electrons[name_u][name_uy].loadChunk(u_y_data, {0u}, {num_to_add}); electrons[name_u][name_uz].loadChunk(u_z_data, {0u}, {num_to_add}); electrons[name_w][name_ww].loadChunk(w_w_data, {0u}, {num_to_add}); + if (m_do_spin_tracking) { + electrons[name_s][name_sx].loadChunk(s_x_data, {0u}, {num_to_add}); + electrons[name_s][name_sy].loadChunk(s_y_data, {0u}, {num_to_add}); + electrons[name_s][name_sz].loadChunk(s_z_data, {0u}, {num_to_add}); + } series.flush(); @@ -1052,6 +1105,7 @@ InitBeamFromFile (const std::string input_file, auto old_size = particle_tile.size(); auto new_size = old_size + num_to_add; particle_tile.resize(new_size); + const auto ptd = particle_tile.getParticleTileData(); const auto enforceBC = EnforceBC(); @@ -1064,7 +1118,11 @@ InitBeamFromFile (const std::string input_file, const input_type * const u_x_ptr = u_x_data.get(); const input_type * const u_y_ptr = u_y_data.get(); const input_type * const u_z_ptr = u_z_data.get(); + const input_type * const s_x_ptr = m_do_spin_tracking ? s_x_data.get() : nullptr; + const input_type * const s_y_ptr = m_do_spin_tracking ? s_y_data.get() : nullptr; + const input_type * const s_z_ptr = m_do_spin_tracking ? s_z_data.get() : nullptr; const input_type * const w_w_ptr = w_w_data.get(); + const bool do_spin_tracking = m_do_spin_tracking; amrex::ParallelFor(amrex::Long(num_to_add), [=] AMREX_GPU_DEVICE (const amrex::Long i) { @@ -1075,8 +1133,11 @@ InitBeamFromFile (const std::string input_file, static_cast(u_x_ptr[i] * unit_ux), // = gamma * beta static_cast(u_y_ptr[i] * unit_uy), static_cast(u_z_ptr[i] * unit_uz), + do_spin_tracking ? static_cast(s_x_ptr[i]) : 0., + do_spin_tracking ? static_cast(s_y_ptr[i]) : 0., + do_spin_tracking ? static_cast(s_z_ptr[i]) : 0., static_cast(w_w_ptr[i] * unit_ww), - pid, i, phys_const.c, enforceBC); + pid, i, phys_const.c, enforceBC, do_spin_tracking); }); amrex::Gpu::streamSynchronize();