Skip to content

Commit

Permalink
Joined Array support. TODO: verify const var dim limit in openpmd-api,
Browse files Browse the repository at this point in the history 
and BTD
  • Loading branch information
@guj
guj committed Sep 3, 2024
1 parent 452ae55 commit 40676dd
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 21 deletions.
5 changes: 4 additions & 1 deletion Source/Diagnostics/WarpXOpenPMD.H
View file
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 @@ -336,6 +336,9 @@ private:
int m_MPIRank = 0;
int m_MPISize = 1;

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

openPMD::IterationEncoding m_Encoding = openPMD::IterationEncoding::fileBased;
std::string m_OpenPMDFileType = "bp"; //! MPI-parallel openPMD backend: bp or h5
std::string m_OpenPMDoptions = "{}"; //! JSON option string for openPMD::Series constructor
Expand Down
84 changes: 64 additions & 20 deletions Source/Diagnostics/WarpXOpenPMD.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,23 @@ WarpXOpenPMDPlot::WarpXOpenPMDPlot (
{
m_OpenPMDoptions = detail::getSeriesOptions(operator_type, operator_parameters,
engine_type, engine_parameters);
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()<<" ... [Using Joined Array for particles in ADIOS??] "<<m_ApplyJoinedArray<<std::endl;
}

openPMD::Extent WarpXOpenPMDPlot::ProperExtent (unsigned long long n, bool init) const
{
if (!m_ApplyJoinedArray)
return {n};

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

WarpXOpenPMDPlot::~WarpXOpenPMDPlot ()
Expand Down Expand Up @@ -529,7 +546,10 @@ 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 (unsigned i = 0, n = particle_diags.size(); i < n; ++i) {

Expand Down Expand Up @@ -665,7 +685,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 @@ -711,11 +731,11 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
SetupRealProperties(pc, currSpecies, write_real_comp, real_comp_names, write_int_comp, int_comp_names,
NewParticleVectorSize, isBTD);
}

#ifdef NO_JOINED_DIM // need to see how to revive const vars in openpmd (needs the size which should get from gathering and we skip gathering with joined array)
if (is_last_flush_to_step) {
SetConstParticleRecordsEDPIC(currSpecies, positionComponents, NewParticleVectorSize, charge, mass);
}

#endif
// open files from all processors, in case some will not contribute below
m_Series->flush();

Expand Down Expand Up @@ -747,6 +767,7 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
} // pti
} // currentLevel

#ifdef NO_JOINED_DIM // possibly not needed, will verify
// work-around for BTD particle resize in ADIOS2
//
// This issues an empty ADIOS2 Put to make sure the new global shape
Expand All @@ -768,7 +789,7 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
[[fallthrough]];
case openPMD::Datatype::DOUBLE : {
auto empty_data = std::make_shared<amrex::ParticleReal>();
comp.storeChunk(empty_data, {uint64_t(0)}, {uint64_t(0)});
comp.storeChunk(empty_data, {uint64_t(0)}, {uint64_t(0)}); // to do
break;
}
case openPMD::Datatype::UINT : {
Expand Down Expand Up @@ -796,7 +817,7 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
}
}
}

#endif // skipping workaround for joind_dim case
m_Series->flush();
}

Expand All @@ -811,8 +832,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 @@ -904,7 +925,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 @@ -935,10 +956,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 @@ -958,7 +979,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 @@ -968,7 +989,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 @@ -984,8 +1005,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 @@ -1003,7 +1024,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 @@ -1477,14 +1498,18 @@ 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()}
{
m_ParticleCounterByLevel.resize(pc->finestLevel()+1);
m_ParticleOffsetAtRank.resize(pc->finestLevel()+1);
m_ParticleSizeAtRank.resize(pc->finestLevel()+1);
WARPX_PROFILE("WarpXParticleCounter");

unsigned long long zero=0;
m_ParticleCounterByLevel.resize(pc->finestLevel()+1, zero);
m_ParticleOffsetAtRank.resize(pc->finestLevel()+1, zero);
m_ParticleSizeAtRank.resize(pc->finestLevel()+1, zero);

if (!doLazyCount) {
for (auto currentLevel = 0; currentLevel <= pc->finestLevel(); currentLevel++)
{
long numParticles = 0; // numParticles in this processor
Expand All @@ -1510,6 +1535,25 @@ WarpXParticleCounter::WarpXParticleCounter (ParticleContainer* pc):

m_Total += sum;
}
} // !doLazyCount
else
{
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;
}

//unsigned long long offset=0; // offset of this level
//unsigned long long sum=0; // numParticles in this level (sum from all processors)

m_Total += numParticles;
}

} // else
}


Expand Down

0 comments on commit 40676dd

Please sign in to comment.