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.

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

Diagnostics: Remove ASCII Particles #443

Merged
merged 2 commits into from
Nov 3, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
1 change: 0 additions & 1 deletion src/particles/diagnostics/DiagnosticOutput.H
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ namespace impactx::diagnostics
*/
enum class OutputType
{
PrintParticles, ///< ASCII diagnostics, for small tests only
PrintNonlinearLensInvariants, ///< ASCII diagnostics for the IOTA nonlinear lens, for small tests only
PrintRefParticle, ///< ASCII diagnostics, for small tests only
ax3l marked this conversation as resolved.
Show resolved Hide resolved
PrintReducedBeamCharacteristics ///< ASCII diagnostics, for beam momenta and Twiss parameters
Expand Down
225 changes: 98 additions & 127 deletions src/particles/diagnostics/DiagnosticOutput.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -41,9 +41,7 @@ namespace impactx::diagnostics

// write file header per MPI RANK
if (!append) {
if (otype == OutputType::PrintParticles) {
file_handler << "id x y t px py pt\n";
} else if (otype == OutputType::PrintNonlinearLensInvariants) {
if (otype == OutputType::PrintNonlinearLensInvariants) {
file_handler << "id H I\n";
} else if (otype == OutputType::PrintRefParticle) {
file_handler << "step s x y z t px py pz pt\n";
Expand All @@ -61,9 +59,29 @@ namespace impactx::diagnostics
}
}

if (otype == OutputType::PrintReducedBeamCharacteristics) {
std::unordered_map<std::string, amrex::ParticleReal> const rbc = diagnostics::reduced_beam_characteristics(
pc);
if (otype == OutputType::PrintRefParticle) {
// preparing to access reference particle data: RefPart
RefPart const ref_part = pc.GetRefParticle();

amrex::ParticleReal const s = ref_part.s;
amrex::ParticleReal const x = ref_part.x;
amrex::ParticleReal const y = ref_part.y;
amrex::ParticleReal const z = ref_part.z;
amrex::ParticleReal const t = ref_part.t;
amrex::ParticleReal const px = ref_part.px;
amrex::ParticleReal const py = ref_part.py;
amrex::ParticleReal const pz = ref_part.pz;
amrex::ParticleReal const pt = ref_part.pt;

// write particle data to file
file_handler
<< step << " " << s << " "
<< x << " " << y << " " << z << " " << t << " "
<< px << " " << py << " " << pz << " " << pt << "\n";
} // if( otype == OutputType::PrintRefParticle)
else if (otype == OutputType::PrintReducedBeamCharacteristics) {
std::unordered_map<std::string, amrex::ParticleReal> const rbc =
diagnostics::reduced_beam_characteristics(pc);

file_handler << step << " " << rbc.at("s") << " " << rbc.at("ref_beta_gamma") << " "
<< rbc.at("x_mean") << " " << rbc.at("y_mean") << " " << rbc.at("t_mean") << " "
Expand All @@ -76,127 +94,80 @@ namespace impactx::diagnostics
<< rbc.at("charge_C") << "\n";
} // if( otype == OutputType::PrintReducedBeamCharacteristics)

// create a host-side particle buffer
// todo: NOT needed for OutputType::PrintRefParticle and OutputType::PrintReducedBeamCharacteristics
auto tmp = pc.make_alike<amrex::PinnedArenaAllocator>();

// copy device-to-host
bool const local = true;
tmp.copyParticles(pc, local);

// loop over refinement levels
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) {
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::px].dataPtr();
amrex::ParticleReal const *const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr();
amrex::ParticleReal const *const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr();

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);
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];
amrex::ParticleReal const pt = part_pt[i];

// write particle data to file
file_handler
<< global_id << " "
<< x << " " << y << " " << t << " "
<< px << " " << py << " " << pt << "\n";
} // i=0...np
} // if( otype == OutputType::PrintParticles)
else if (otype == OutputType::PrintNonlinearLensInvariants) {

using namespace amrex::literals;

// Parse the diagnostic parameters
amrex::ParmParse pp_diag("diag");

amrex::ParticleReal alpha = 0.0;
pp_diag.queryAdd("alpha", alpha);

amrex::ParticleReal beta = 1.0;
pp_diag.queryAdd("beta", beta);

amrex::ParticleReal tn = 0.4;
pp_diag.queryAdd("tn", tn);

amrex::ParticleReal cn = 0.01;
pp_diag.queryAdd("cn", cn);

NonlinearLensInvariants const nonlinear_lens_invariants(alpha, beta, tn, cn);

// 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);
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];

// calculate invariants of motion
NonlinearLensInvariants::Data const HI_out =
nonlinear_lens_invariants(x, y, px, py);

// write particle invariant data to file
file_handler
<< global_id << " "
<< HI_out.H << " " << HI_out.I << "\n";

} // i=0...np
} // if( otype == OutputType::PrintInvariants)
if (otype == OutputType::PrintRefParticle) {
// print reference particle to file

// preparing to access reference particle data: RefPart
RefPart const ref_part = pc.GetRefParticle();

amrex::ParticleReal const s = ref_part.s;
amrex::ParticleReal const x = ref_part.x;
amrex::ParticleReal const y = ref_part.y;
amrex::ParticleReal const z = ref_part.z;
amrex::ParticleReal const t = ref_part.t;
amrex::ParticleReal const px = ref_part.px;
amrex::ParticleReal const py = ref_part.py;
amrex::ParticleReal const pz = ref_part.pz;
amrex::ParticleReal const pt = ref_part.pt;

// write particle data to file
file_handler
<< step << " " << s << " "
<< x << " " << y << " " << z << " " << t << " "
<< px << " " << py << " " << pz << " " << pt << "\n";
} // if( otype == OutputType::PrintRefParticle)
} // end loop over all particle boxes
} // env mesh-refinement level loop
// TODO: add as an option to the monitor element
if (otype == OutputType::PrintNonlinearLensInvariants) {
// create a host-side particle buffer
auto tmp = pc.make_alike<amrex::PinnedArenaAllocator>();

// copy all particles from device to host
bool const local = true;
tmp.copyParticles(pc, local);

// loop over refinement levels
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) {
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::px].dataPtr();
amrex::ParticleReal const *const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr();

if (otype == OutputType::PrintNonlinearLensInvariants) {
using namespace amrex::literals;

// Parse the diagnostic parameters
amrex::ParmParse pp_diag("diag");

amrex::ParticleReal alpha = 0.0;
pp_diag.queryAdd("alpha", alpha);

amrex::ParticleReal beta = 1.0;
pp_diag.queryAdd("beta", beta);

amrex::ParticleReal tn = 0.4;
pp_diag.queryAdd("tn", tn);

amrex::ParticleReal cn = 0.01;
pp_diag.queryAdd("cn", cn);

NonlinearLensInvariants const nonlinear_lens_invariants(alpha, beta, tn, cn);

// print out particles
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);
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];

// calculate invariants of motion
NonlinearLensInvariants::Data const HI_out =
nonlinear_lens_invariants(x, y, px, py);

// write particle invariant data to file
file_handler
<< global_id << " "
<< HI_out.H << " " << HI_out.I << "\n";

} // i=0...np
} // if( otype == OutputType::PrintInvariants)
} // end loop over all particle boxes
} // env mesh-refinement level loop
}
}

} // namespace impactx::diagnostics