Skip to content

Joined Array support is added for ADIOS when writing out particles. #5804

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

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

Joined Array support is added for ADIOS when writing out particles. #5804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5b46b52
Joined Array support added for ADIOS
guj Mar 28, 2025
25c08a6
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 28, 2025
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
5 changes: 4 additions & 1 deletion Source/Diagnostics/WarpXOpenPMD.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ public:
using ParticleContainer = typename WarpXParticleContainer::ContainerLike<amrex::PinnedArenaAllocator>;
using ParticleIter = typename amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>;

WarpXParticleCounter (ParticleContainer* pc);
WarpXParticleCounter (ParticleContainer* pc, bool doLazyCount);
[[nodiscard]] unsigned long GetTotalNumParticles () const {return m_Total;}

std::vector<unsigned long long> m_ParticleOffsetAtRank;
Expand Down Expand Up @@ -349,6 +349,9 @@ private:
int m_MPIRank = 0;
int m_MPISize = 1;

bool m_ApplyJoinedArray = false; // joined array
[[nodiscard]] openPMD::Extent ProperExtent(unsigned long long n, [[maybe_unused]] bool init) const;

openPMD::IterationEncoding m_Encoding = openPMD::IterationEncoding::fileBased;
std::string m_OpenPMDFileType = "bp5"; //! MPI-parallel openPMD backend: bp5, bp4 or h5
std::string m_OpenPMDoptions = "{}"; //! JSON option string for openPMD::Series constructor
Expand Down
72 changes: 59 additions & 13 deletions Source/Diagnostics/WarpXOpenPMD.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -390,6 +390,32 @@ WarpXOpenPMDPlot::WarpXOpenPMDPlot (
{
m_OpenPMDoptions = ::detail::getSeriesOptions(operator_type, operator_parameters,
engine_type, engine_parameters);

#if OPENPMDAPI_VERSION_GE(0, 16, 0)
// Automatically use Joined Array in particles for bp4/bp5 files
// The joined dimention feature is supported in openPMD-api 0.16+
std::size_t found_bp4 = m_OpenPMDFileType.find("bp4");
std::size_t found_bp5 = m_OpenPMDFileType.find("bp5");
if ( (found_bp4 != std::string::npos) || (found_bp5 != std::string::npos) ) {
m_ApplyJoinedArray = true;
amrex::Print()<< Utils::TextMsg::Info(" [Joined Array will be applied for particles in ADIOS] ");
}
#endif
}

[[nodiscard]] openPMD::Extent WarpXOpenPMDPlot::ProperExtent (unsigned long long n, [[maybe_unused]] bool init) const
{
#if OPENPMDAPI_VERSION_GE(0, 16, 0)
if (!m_ApplyJoinedArray)
return {n};

if (init)
return {openPMD::Dataset::JOINED_DIMENSION};
else
return {};
#else
return {n};
#endif
}

WarpXOpenPMDPlot::~WarpXOpenPMDPlot ()
Expand Down Expand Up @@ -546,7 +572,11 @@ WarpXOpenPMDPlot::WriteOpenPMDParticles (const amrex::Vector<ParticleDiag>& part
const bool isLastBTDFlush
)
{
WARPX_PROFILE("WarpXOpenPMDPlot::WriteOpenPMDParticles()");
std::string name="WarpXOpenPMDPlot::WriteOpenPMDParticles()";
if ( m_ApplyJoinedArray ) {
name +="_joined ";
}
WARPX_PROFILE(name);

for (const auto & particle_diag : particle_diags) {
WarpXParticleContainer* pc = particle_diag.getParticleContainer();
Expand Down Expand Up @@ -702,7 +732,7 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc->NumRealComps());
AMREX_ALWAYS_ASSERT(int_comp_names.size() == pc->NumIntComps());

WarpXParticleCounter counter(pc);
WarpXParticleCounter counter(pc, m_ApplyJoinedArray);
auto const num_dump_particles = counter.GetTotalNumParticles();

openPMD::Iteration currIteration = GetIteration(iteration, isBTD);
Expand Down Expand Up @@ -848,8 +878,8 @@ WarpXOpenPMDPlot::SetupRealProperties (ParticleContainer const * pc,
{
std::string options = "{}";
if (isBTD) { options = "{ \"resizable\": true }"; }
auto dtype_real = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np}, options);
auto dtype_int = openPMD::Dataset(openPMD::determineDatatype<int>(), {np}, options);
auto dtype_real = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), ProperExtent(np, true), options);
auto dtype_int = openPMD::Dataset(openPMD::determineDatatype<int>(), ProperExtent(np, true), options);
//
// the beam/input3d showed write_real_comp.size() = 16 while only 10 real comp names
// so using the min to be safe.
Expand Down Expand Up @@ -941,7 +971,7 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
{
// todo: add support to not write the particle index
getComponentRecord("id").storeChunkRaw(
soa.GetIdCPUData().data(), {offset}, {numParticleOnTile64});
soa.GetIdCPUData().data(), ProperExtent(offset, false), {numParticleOnTile64});
}

// here we the save the SoA properties (real)
Expand Down Expand Up @@ -972,10 +1002,10 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
if (write_real_comp[1]) { y.get()[i] = yp; }
}
if (write_real_comp[0]) {
getComponentRecord(real_comp_names[0]).storeChunk(x, {offset}, {numParticleOnTile64});
getComponentRecord(real_comp_names[0]).storeChunk(x, ProperExtent(offset, false), {numParticleOnTile64});
}
if (write_real_comp[1]) {
getComponentRecord(real_comp_names[1]).storeChunk(y, {offset}, {numParticleOnTile64});
getComponentRecord(real_comp_names[1]).storeChunk(y, ProperExtent(offset, false), {numParticleOnTile64});
}
#endif

Expand All @@ -995,7 +1025,7 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
#endif
if (write_real_comp[idx]) {
getComponentRecord(real_comp_names[idx]).storeChunkRaw(
soa.GetRealData(soa_r_idx).data(), {offset}, {numParticleOnTile64});
soa.GetRealData(soa_r_idx).data(), ProperExtent(offset, false), {numParticleOnTile64});
}
}
}
Expand All @@ -1005,7 +1035,7 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
for (auto idx=0; idx<int_counter; idx++) {
if (write_int_comp[idx]) {
getComponentRecord(int_comp_names[idx]).storeChunkRaw(
soa.GetIntData(idx).data(), {offset}, {numParticleOnTile64});
soa.GetIntData(idx).data(), ProperExtent(offset, false), {numParticleOnTile64});
}
}
}
Expand All @@ -1021,8 +1051,8 @@ WarpXOpenPMDPlot::SetupPos (
{
std::string options = "{}";
if (isBTD) { options = "{ \"resizable\": true }"; }
auto realType = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np}, options);
auto idType = openPMD::Dataset(openPMD::determineDatatype< uint64_t >(), {np}, options);
auto realType = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), ProperExtent(np, true), options);
auto idType = openPMD::Dataset(openPMD::determineDatatype< uint64_t >(), ProperExtent(np, true), options);

for( auto const& comp : positionComponents ) {
currSpecies["position"][comp].resetDataset( realType );
Expand All @@ -1040,7 +1070,7 @@ WarpXOpenPMDPlot::SetConstParticleRecordsEDPIC (
amrex::ParticleReal const charge,
amrex::ParticleReal const mass)
{
auto realType = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np});
auto realType = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), ProperExtent(np, true));
const auto *const scalar = openPMD::RecordComponent::SCALAR;

// define record shape to be number of particles
Expand Down Expand Up @@ -1516,7 +1546,7 @@ WarpXOpenPMDPlot::WriteOpenPMDFieldsAll ( //const std::string& filename,
//
//
//
WarpXParticleCounter::WarpXParticleCounter (ParticleContainer* pc):
WarpXParticleCounter::WarpXParticleCounter (ParticleContainer* pc, bool doLazyCount):
m_MPIRank{amrex::ParallelDescriptor::MyProc()},
m_MPISize{amrex::ParallelDescriptor::NProcs()}
{
Expand All @@ -1525,6 +1555,21 @@ WarpXParticleCounter::WarpXParticleCounter (ParticleContainer* pc):
m_ParticleOffsetAtRank.resize(pc->finestLevel()+1);
m_ParticleSizeAtRank.resize(pc->finestLevel()+1);

if (doLazyCount) {
for (auto currentLevel = 0; currentLevel <= pc->finestLevel(); currentLevel++)
{
long numParticles = 0; // numParticles in this processor
for (ParticleIter pti(*pc, currentLevel); pti.isValid(); ++pti) {
auto numParticleOnTile = pti.numParticles();
numParticles += numParticleOnTile;
}
m_Total += numParticles;
}
return;
}

{// MPI gather to get offsets from each rank

for (auto currentLevel = 0; currentLevel <= pc->finestLevel(); currentLevel++)
{
long numParticles = 0; // numParticles in this processor
Expand All @@ -1550,6 +1595,7 @@ WarpXParticleCounter::WarpXParticleCounter (ParticleContainer* pc):

m_Total += sum;
}
}
}


Expand Down
Loading