Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Beam read_from_file also reads spin #1172

Open
wants to merge 9 commits into
base: development
Choose a base branch
from
Open

Beam read_from_file also reads spin #1172

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
04dec14
spin polarized read from file, dirty atm
MaxThevenet Oct 7, 2024
f545ac1
do some cleaning
MaxThevenet Oct 7, 2024
2c8e21a
cleaning and make sure still works when openpmd file has no spin data
MaxThevenet Oct 7, 2024
3374191
fix merge conflicts
MaxThevenet Oct 14, 2024
ccebfa4
fix Arena
MaxThevenet Oct 14, 2024
573b319
move define in case of spin
MaxThevenet Oct 14, 2024
39daa8e
add initslice
MaxThevenet Oct 14, 2024
2d1975b
move init slice
MaxThevenet Oct 14, 2024
d92f107
Update src/particles/beam/BeamParticleContainerInit.cpp
MaxThevenet Oct 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 19 additions & 12 deletions src/particles/beam/BeamParticleContainer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,18 @@ 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);
}
AlexanderSinn marked this conversation as resolved.
Show resolved Hide resolved
}
auto& soa = getBeamInitSlice().GetStructOfArrays();
soa.GetIdCPUData().setArena(
m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena());
Expand All @@ -102,15 +114,6 @@ 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);
for (auto& beam_tile : m_slices) {
// Use 3 real and 0 int runtime components
beam_tile.define(3, 0);
}
}
}

amrex::Real
Expand Down Expand Up @@ -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) {
Expand All @@ -379,20 +383,23 @@ 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][ip];
ptd.m_runtime_rdata[1][ip] = ptd_init.m_runtime_rdata[1][ip];
ptd.m_runtime_rdata[2][ip] = ptd_init.m_runtime_rdata[2][ip];
}
ptd.idcpu(ip) = ptd_init.idcpu(idx_src);
ptd.idata(BeamIdx::nsubcycles)[ip] = 0;
ptd.idata(BeamIdx::mr_level)[ip] = 0;
}
);
}

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];
Expand Down
76 changes: 70 additions & 6 deletions src/particles/beam/BeamParticleContainerInit.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 ="";
AlexanderSinn marked this conversation as resolved.
Show resolved Hide resolved
std::string name_m ="", name_mm ="";
std::string name_q ="", name_qq ="";
std::string name_g ="", name_gg ="";
Expand Down Expand Up @@ -873,6 +881,24 @@ InitBeamFromFile (const std::string input_file,
}
}
}
else if(units == std::array<double,7> {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;
}
}
}
}
}
}
}
Expand Down Expand Up @@ -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];
Expand All @@ -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<input_type> u_z_data{ reinterpret_cast<input_type*>(
amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del};
std::shared_ptr<input_type> s_x_data{ reinterpret_cast<input_type*>(
amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ?
sizeof(input_type)*num_to_add : 0) ), del};
std::shared_ptr<input_type> s_y_data{ reinterpret_cast<input_type*>(
amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ?
sizeof(input_type)*num_to_add : 0) ), del};
std::shared_ptr<input_type> s_z_data{ reinterpret_cast<input_type*>(
amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ?
sizeof(input_type)*num_to_add: 0 ) ), del};
std::shared_ptr<input_type> w_w_data{ reinterpret_cast<input_type*>(
amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del};

Expand All @@ -997,6 +1045,11 @@ InitBeamFromFile (const std::string input_file,
electrons[name_u][name_uy].loadChunk<input_type>(u_y_data, {0u}, {num_to_add});
electrons[name_u][name_uz].loadChunk<input_type>(u_z_data, {0u}, {num_to_add});
electrons[name_w][name_ww].loadChunk<input_type>(w_w_data, {0u}, {num_to_add});
if (m_do_spin_tracking) {
electrons[name_s][name_sx].loadChunk<input_type>(s_x_data, {0u}, {num_to_add});
electrons[name_s][name_sy].loadChunk<input_type>(s_y_data, {0u}, {num_to_add});
electrons[name_s][name_sz].loadChunk<input_type>(s_z_data, {0u}, {num_to_add});
}

series.flush();

Expand Down Expand Up @@ -1051,7 +1104,11 @@ InitBeamFromFile (const std::string input_file,
auto& particle_tile = getBeamInitSlice();
auto old_size = particle_tile.size();
auto new_size = old_size + num_to_add;
if (m_do_spin_tracking) {
particle_tile.define(3, 0);
}
MaxThevenet marked this conversation as resolved.
Show resolved Hide resolved
particle_tile.resize(new_size);

const auto ptd = particle_tile.getParticleTileData();
const auto enforceBC = EnforceBC();

Expand All @@ -1064,7 +1121,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) {
Expand All @@ -1075,8 +1136,11 @@ InitBeamFromFile (const std::string input_file,
static_cast<amrex::Real>(u_x_ptr[i] * unit_ux), // = gamma * beta
static_cast<amrex::Real>(u_y_ptr[i] * unit_uy),
static_cast<amrex::Real>(u_z_ptr[i] * unit_uz),
do_spin_tracking ? static_cast<amrex::Real>(s_x_ptr[i]) : 0.,
do_spin_tracking ? static_cast<amrex::Real>(s_y_ptr[i]) : 0.,
do_spin_tracking ? static_cast<amrex::Real>(s_z_ptr[i]) : 0.,
static_cast<amrex::Real>(w_w_ptr[i] * unit_ww),
pid, i, phys_const.c, enforceBC);
pid, i, phys_const.c, enforceBC, do_spin_tracking);
});

amrex::Gpu::streamSynchronize();
Expand Down
Loading