Skip to content

Commit

Permalink
relax constraint that real_comp_names.size() == pc.NumRealComps() + N…
Browse files Browse the repository at this point in the history 
…StructReal for pure SoA plotfiles (AMReX-Codes#3717)

The proposed changes:
- [x] fix a bug or incorrect behavior in AMReX
- [ ] add new capabilities to AMReX
- [ ] changes answers in the test suite to more than roundoff level
- [ ] are likely to significantly affect the results of downstream AMReX
users
- [ ] include documentation in the code and/or rst files, if appropriate
  • Loading branch information
@atmyers
atmyers authored Jan 23, 2024
1 parent 2fc8a34 commit 73b2155
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 29 deletions.
74 changes: 56 additions & 18 deletions Src/Particle/AMReX_ParticleIO.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
{
Vector<int> write_real_comp;
Vector<std::string> tmp_real_comp_names;
for (int i = 0; i < NStructReal + NumRealComps(); ++i )
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();

for (int i = 0; i < nrc; ++i )
{
write_real_comp.push_back(1);
if (real_comp_names.empty())
Expand Down Expand Up @@ -96,7 +98,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
{
Vector<int> write_real_comp;
Vector<std::string> real_comp_names;
for (int i = 0; i < NStructReal + NumRealComps(); ++i )
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();

for (int i = 0; i < nrc; ++i )
{
write_real_comp.push_back(1);
std::stringstream ss;
Expand Down Expand Up @@ -127,11 +131,16 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
const Vector<std::string>& real_comp_names,
const Vector<std::string>& int_comp_names) const
{
AMREX_ASSERT(real_comp_names.size() == NStructReal + NumRealComps());
if constexpr(ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal);
}
AMREX_ASSERT( int_comp_names.size() == NStructInt + NumIntComps() );

Vector<int> write_real_comp;
for (int i = 0; i < NStructReal + NumRealComps(); ++i) {
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
for (int i = 0; i < nrc; ++i) {
write_real_comp.push_back(1);
}

Expand All @@ -153,10 +162,15 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
::WritePlotFile (const std::string& dir, const std::string& name,
const Vector<std::string>& real_comp_names) const
{
AMREX_ASSERT(real_comp_names.size() == NStructReal + NumRealComps());
if constexpr(ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal);
}

Vector<int> write_real_comp;
for (int i = 0; i < NStructReal + NumRealComps(); ++i) {
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
for (int i = 0; i < nrc; ++i) {
write_real_comp.push_back(1);
}

Expand Down Expand Up @@ -188,11 +202,17 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
const Vector<int>& write_real_comp,
const Vector<int>& write_int_comp) const
{
AMREX_ASSERT(write_real_comp.size() == NStructReal + NumRealComps());

if constexpr(ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(write_real_comp.size() == NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(write_real_comp.size() == NumRealComps() + NStructReal);
}
AMREX_ASSERT(write_int_comp.size() == NStructInt + NArrayInt );

Vector<std::string> real_comp_names;
for (int i = 0; i < NStructReal + NumRealComps(); ++i )
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
for (int i = 0; i < nrc; ++i )
{
std::stringstream ss;
ss << "real_comp" << i;
Expand Down Expand Up @@ -239,7 +259,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
{
Vector<int> write_real_comp;
Vector<std::string> real_comp_names;
for (int i = 0; i < NStructReal + NumRealComps(); ++i )
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();

for (int i = 0; i < nrc; ++i )
{
write_real_comp.push_back(1);
std::stringstream ss;
Expand Down Expand Up @@ -271,11 +293,16 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
const Vector<std::string>& real_comp_names,
const Vector<std::string>& int_comp_names, F&& f) const
{
AMREX_ASSERT(real_comp_names.size() == NStructReal + NumRealComps());
if constexpr(ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal);
}
AMREX_ASSERT( int_comp_names.size() == NStructInt + NArrayInt );

Vector<int> write_real_comp;
for (int i = 0; i < NStructReal + NumRealComps(); ++i) {
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
for (int i = 0; i < nrc; ++i) {
write_real_comp.push_back(1);
}

Expand All @@ -298,10 +325,15 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
::WritePlotFile (const std::string& dir, const std::string& name,
const Vector<std::string>& real_comp_names, F&& f) const
{
AMREX_ASSERT(real_comp_names.size() == NStructReal + NumRealComps());
if constexpr(ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == NumRealComps() + NStructReal);
}

Vector<int> write_real_comp;
for (int i = 0; i < NStructReal + NumRealComps(); ++i) {
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
for (int i = 0; i < nrc; ++i) {
write_real_comp.push_back(1);
}

Expand Down Expand Up @@ -334,11 +366,16 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
const Vector<int>& write_real_comp,
const Vector<int>& write_int_comp, F&& f) const
{
AMREX_ASSERT(write_real_comp.size() == NStructReal + NumRealComps());
if constexpr(ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(write_real_comp.size() == NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(write_real_comp.size() == NumRealComps() + NStructReal);
}
AMREX_ASSERT(write_int_comp.size() == NStructInt + NumIntComps() );

Vector<std::string> real_comp_names;
for (int i = 0; i < NStructReal + NumRealComps(); ++i )
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
for (int i = 0; i < nrc; ++i )
{
std::stringstream ss;
ss << "real_comp" << i;
Expand Down Expand Up @@ -680,8 +717,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig

int nr;
HdrFile >> nr;
if (nr != NStructReal + NumRealComps()) {
amrex::Abort("ParticleContainer::Restart(): nr != NStructReal + NumRealComps()");
int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
if (nr != nrc) {
amrex::Abort("ParticleContainer::Restart(): nr not the expected value");
}

std::string comp_name;
Expand Down Expand Up @@ -929,7 +967,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
readIntData(istuff.dataPtr(), istuff.size(), ifs, FPC::NativeIntDescriptor());

// Then the real data in binary.
const int rChunkSize = AMREX_SPACEDIM + NStructReal + NumRealComps();
const int rChunkSize = ParticleType::is_soa_particle ? NStructReal + NumRealComps() : AMREX_SPACEDIM + NStructReal + NumRealComps();
Vector<RTYPE> rstuff(std::size_t(cnt)*rChunkSize);
ReadParticleRealData(rstuff.dataPtr(), rstuff.size(), ifs);

Expand Down
33 changes: 23 additions & 10 deletions Src/Particle/AMReX_WriteBinaryParticleData.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,7 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
idata.resize(np*iChunkSize);

int num_output_real = 0;
for (int i = 0; i < pc.NumRealComps() + PC::NStructReal; ++i) {
for (int i = 0; i < (int) write_real_comp.size(); ++i) {
if (write_real_comp[i]) { ++num_output_real; }
}

Expand Down Expand Up @@ -361,7 +361,9 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
// extra SoA Real components
const int real_start_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
for (int j = real_start_offset; j < pc.NumRealComps(); j++) {
if (write_real_comp[PC::NStructReal+j]) {
const int write_comp_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
const int write_comp_index = PC::NStructReal+j-write_comp_offset;
if (write_real_comp[write_comp_index]) {
*rptr = (ParticleReal) soa.GetRealData(j)[pindex];
++rptr;
}
Expand Down Expand Up @@ -393,7 +395,11 @@ void WriteBinaryParticleDataSync (PC const& pc,
const int NProcs = ParallelDescriptor::NProcs();
const int IOProcNumber = ParallelDescriptor::IOProcessorNumber();

AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc.NumRealComps() + NStructReal);
if constexpr(PC::ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc.NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc.NumRealComps() + NStructReal);
}
AMREX_ALWAYS_ASSERT( int_comp_names.size() == pc.NumIntComps() + NStructInt);

std::string pdir = dir;
Expand Down Expand Up @@ -476,7 +482,7 @@ void WriteBinaryParticleDataSync (PC const& pc,
}

int num_output_real = 0;
for (int i = 0; i < pc.NumRealComps() + NStructReal; ++i) {
for (int i = 0; i < (int) write_real_comp.size(); ++i) {
if (write_real_comp[i]) { ++num_output_real; }
}

Expand All @@ -492,7 +498,7 @@ void WriteBinaryParticleDataSync (PC const& pc,
HdrFile << num_output_real << '\n';

// Real component names
for (int i = 0; i < NStructReal + pc.NumRealComps(); ++i ) {
for (int i = 0; i < (int) real_comp_names.size(); ++i ) {
if (write_real_comp[i]) { HdrFile << real_comp_names[i] << '\n'; }
}

Expand Down Expand Up @@ -683,7 +689,11 @@ void WriteBinaryParticleDataAsync (PC const& pc,
const int NProcs = ParallelDescriptor::NProcs();
const int IOProcNumber = NProcs - 1;

AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc.NumRealComps() + NStructReal);
if constexpr(PC::ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc.NumRealComps() + NStructReal - AMREX_SPACEDIM); // pure SoA: skip positions
} else {
AMREX_ALWAYS_ASSERT(real_comp_names.size() == pc.NumRealComps() + NStructReal);
}
AMREX_ALWAYS_ASSERT( int_comp_names.size() == pc.NumIntComps() + NStructInt);

Vector<LayoutData<Long> > np_per_grid_local(pc.finestLevel()+1);
Expand Down Expand Up @@ -847,6 +857,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,

int nrc = pc.NumRealComps();
int nic = pc.NumIntComps();
int rnames_size = (int) real_comp_names.size();

auto RD = pc.ParticleRealDescriptor;

Expand Down Expand Up @@ -881,7 +892,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,
}

int num_output_real = 0;
for (int i = 0; i < nrc + NStructReal; ++i) {
for (int i = 0; i < rnames_size; ++i) {
if (write_real_comp[i]) { ++num_output_real; }
}

Expand All @@ -897,7 +908,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,
HdrFile << num_output_real << '\n';

// Real component names
for (int i = 0; i < NStructReal + nrc; ++i ) {
for (int i = 0; i < rnames_size; ++i ) {
if (write_real_comp[i]) { HdrFile << real_comp_names[i] << '\n'; }
}

Expand Down Expand Up @@ -1047,7 +1058,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,

// Write the Real data in binary.
int num_output_real = 0;
for (int i = 0; i < nrc + NStructReal; ++i) {
for (int i = 0; i < rnames_size; ++i) {
if (write_real_comp[i]) { ++num_output_real; }
}

Expand Down Expand Up @@ -1093,7 +1104,9 @@ void WriteBinaryParticleDataAsync (PC const& pc,
const int real_start_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: positions
for (int j = real_start_offset; j < nrc; j++)
{
if (write_real_comp[NStructReal+j])
const int write_comp_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
const int write_comp_index = PC::NStructReal+j-write_comp_offset;
if (write_real_comp[write_comp_index])
{
*rptr = (typename PC::ParticleType::RealType) soa.GetRealData(j)[pindex];
++rptr;
Expand Down
2 changes: 1 addition & 1 deletion Tests/Particles/CheckpointRestartSOA/main.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ void test ()
amrex::Print() << " done \n";

Vector<std::string> particle_realnames;
for (int i = 0; i < NReal; ++i) {
for (int i = 0; i < NReal-AMREX_SPACEDIM; ++i) {
particle_realnames.push_back("particle_real_component_" + std::to_string(i));
}

Expand Down

0 comments on commit 73b2155

Please sign in to comment.