Skip to content

Commit

Permalink
openPMD: Support for mesh refinement (#1842)
Browse files Browse the repository at this point in the history 
* Added support for mesh refinement

* openPMD MR: Cleanup

* Fix BTD

Move un-related streaming API changes out of this PR.

Co-authored-by: Axel Huebl <[email protected]>
  • Loading branch information
@guj @ax3l
guj and ax3l authored May 20, 2021
1 parent 01a4377 commit 45e6483
Show file tree
Hide file tree
Showing 3 changed files with 163 additions and 99 deletions.
4 changes: 2 additions & 2 deletions Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@ FlushFormatOpenPMD::WriteToFile (
m_OpenPMDPlotWriter->SetStep(output_iteration, prefix, isBTD);

// fields: only dumped for coarse level
m_OpenPMDPlotWriter->WriteOpenPMDFields(
varnames, mf[0], geom[0], output_iteration, time, isBTD, full_BTD_snapshot);
m_OpenPMDPlotWriter->WriteOpenPMDFieldsAll(
varnames, mf, geom, output_iteration, time, isBTD, full_BTD_snapshot);

// particles: all (reside only on locally finest level)
m_OpenPMDPlotWriter->WriteOpenPMDParticles(particle_diags);
Expand Down
22 changes: 19 additions & 3 deletions Source/Diagnostics/WarpXOpenPMD.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,10 +117,10 @@ public:

void WriteOpenPMDParticles (const amrex::Vector<ParticleDiag>& particle_diags);

void WriteOpenPMDFields (
void WriteOpenPMDFieldsAll (
const std::vector<std::string>& varnames,
const amrex::MultiFab& mf,
const amrex::Geometry& geom,
const amrex::Vector<amrex::MultiFab>& mf,
amrex::Vector<amrex::Geometry>& geom,
const int iteration, const double time,
bool isBTD = false,
const amrex::Geometry& full_BTD_snapshot=amrex::Geometry() ) const;
Expand All @@ -129,6 +129,22 @@ public:
private:
void Init (openPMD::Access access, bool isBTD);

/** This function does initial setup for the fields when interation is newly created
* @param[in] meshes The meshes in a series
* @param[in] full_geom The geometry
*/
void SetupFields( openPMD::Container< openPMD::Mesh >& meshes, amrex::Geometry& full_geom ) const;

void SetupMeshComp( openPMD::Mesh& mesh,
amrex::Geometry& full_geom,
openPMD::MeshRecordComponent& mesh_comp
) const;

void GetMeshCompNames( int meshLevel,
const std::string& varname,
std::string& field_name,
std::string& comp_name ) const;

/** This function sets up the entries for storing the particle positions, global IDs, and constant records (charge, mass)
*
* @param[in] currSpecies Corresponding openPMD species
Expand Down
236 changes: 142 additions & 94 deletions Source/Diagnostics/WarpXOpenPMD.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,7 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
WarpXParticleCounter counter(pc);

openPMD::Iteration currIteration = m_Series->iterations[iteration];

openPMD::ParticleSpecies currSpecies = currIteration.particles[name];
// meta data for ED-PIC extension
currSpecies.setAttribute( "particleShape", double( WarpX::noz ) );
Expand Down Expand Up @@ -766,56 +767,17 @@ WarpXOpenPMDPlot::SetupPos(
}


//
// this is originally copied from FieldIO.cpp
//
/*
* Set up paramter for mesh container using the geometry (from level 0)
*
* @param [IN] meshes: openPMD-api mesh container
* @param [IN] full_geom: field geometry
*
*/
void
WarpXOpenPMDPlot::WriteOpenPMDFields ( //const std::string& filename,
const std::vector<std::string>& varnames,
const amrex::MultiFab& mf,
const amrex::Geometry& geom, // geometry of the mf/Fab
const int iteration,
const double time, bool isBTD,
const amrex::Geometry& full_BTD_snapshot ) const
WarpXOpenPMDPlot::SetupFields( openPMD::Container< openPMD::Mesh >& meshes,
amrex::Geometry& full_geom ) const
{
//This is AMReX's tiny profiler. Possibly will apply it later
WARPX_PROFILE("WarpXOpenPMDPlot::WriteOpenPMDFields()");

AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_Series != nullptr, "openPMD series must be initialized");

amrex::Geometry full_geom = geom;
if( isBTD )
full_geom = full_BTD_snapshot;

// is this either a regular write (true) or the first write in a
// backtransformed diagnostic (BTD):
bool const first_write_to_iteration = ! m_Series->iterations.contains( iteration );

int const ncomp = mf.nComp();

// Create a few vectors that store info on the global domain
// Swap the indices for each of them, since AMReX data is Fortran order
// and since the openPMD API assumes contiguous C order
// - Size of the box, in integer number of cells
amrex::Box const & global_box = full_geom.Domain();
auto const global_size = getReversedVec(global_box.size());
// - Grid spacing
std::vector<double> const grid_spacing = getReversedVec(full_geom.CellSize());
// - Global offset
std::vector<double> const global_offset = getReversedVec(full_geom.ProbLo());
// - AxisLabels
std::vector<std::string> axis_labels = detail::getFieldAxisLabels();

// Prepare the type of dataset that will be written
openPMD::Datatype const datatype = openPMD::determineDatatype<amrex::Real>();
auto const dataset = openPMD::Dataset(datatype, global_size);

// meta data
auto series_iteration = m_Series->iterations[iteration];
auto meshes = series_iteration.meshes;
if( first_write_to_iteration ) {
series_iteration.setTime( time );

// meta data for ED-PIC extension
auto const period = full_geom.periodicity(); // TODO double-check: is this the proper global bound or of some level?
std::vector<std::string> fieldBoundary(6, "reflecting");
Expand Down Expand Up @@ -875,16 +837,57 @@ WarpXOpenPMDPlot::WriteOpenPMDFields ( //const std::string& filename,
}());
if (WarpX::do_dive_cleaning)
meshes.setAttribute("chargeCorrectionParameters", "period=1");
}
}


// Loop through the different components, i.e. different fields stored in mf
for (int icomp=0; icomp<ncomp; icomp++){
std::string const & varname = varnames[icomp];
/*
* Setup component properties for a field mesh
* @param [IN]: mesh a mesh field
* @param [IN]: full_geom geometry for the mesh
* @param [IN]: mesh_comp a component for the mesh
*/
void
WarpXOpenPMDPlot::SetupMeshComp( openPMD::Mesh& mesh,
amrex::Geometry& full_geom,
openPMD::MeshRecordComponent& mesh_comp ) const
{
amrex::Box const & global_box = full_geom.Domain();
auto const global_size = getReversedVec(global_box.size());
// - Grid spacing
std::vector<double> const grid_spacing = getReversedVec(full_geom.CellSize());
// - Global offset
std::vector<double> const global_offset = getReversedVec(full_geom.ProbLo());
// - AxisLabels
std::vector<std::string> axis_labels = detail::getFieldAxisLabels();

// Prepare the type of dataset that will be written
openPMD::Datatype const datatype = openPMD::determineDatatype<amrex::Real>();
auto const dataset = openPMD::Dataset(datatype, global_size);

mesh.setDataOrder(openPMD::Mesh::DataOrder::C);
mesh.setAxisLabels(axis_labels);
mesh.setGridSpacing(grid_spacing);
mesh.setGridGlobalOffset(global_offset);
mesh.setAttribute("fieldSmoothing", "none");
mesh_comp.resetDataset(dataset);

// assume fields are scalar unless they match the following match of known vector fields
std::string field_name = varname;
std::string comp_name = openPMD::MeshRecordComponent::SCALAR;
}

/*
* Get component names of a field for openPMD-api book-keeping
* Level is reflected as _lvl<meshLevel>
*
* @param meshLevel [IN]: level of mesh
* @param varname [IN]: name from WarpX
* @param field_name [OUT]: field name for openPMD-api output
* @param comp_name [OUT]: comp name for openPMD-api output
*/
void
WarpXOpenPMDPlot::GetMeshCompNames( int meshLevel,
const std::string& varname,
std::string& field_name,
std::string& comp_name ) const
{
if (varname.size() >= 2u ) {
std::string const varname_1st = varname.substr(0u, 1u); // 1st character
std::string const varname_2nd = varname.substr(1u, 1u); // 2nd character
Expand All @@ -904,51 +907,96 @@ WarpXOpenPMDPlot::WriteOpenPMDFields ( //const std::string& filename,
}
}

// Setup the mesh record accordingly
auto mesh = meshes[field_name];
// MultiFab: Fortran order of indices and axes;
// we invert (only) meta-data arrays to assign labels and offsets in the
// order: slowest to fastest varying index when accessing the mesh
// contiguously (as 1D flattened logical memory)
if( first_write_to_iteration ) {
mesh.setDataOrder(openPMD::Mesh::DataOrder::C);
mesh.setAxisLabels(axis_labels);
mesh.setGridSpacing(grid_spacing);
mesh.setGridGlobalOffset(global_offset);
mesh.setAttribute("fieldSmoothing", "none");
detail::setOpenPMDUnit(mesh, field_name);
}
if ( 0 == meshLevel )
return;

// Create a new mesh record component, and store the associated metadata
auto mesh_comp = mesh[comp_name];
if( first_write_to_iteration ) {
mesh_comp.resetDataset(dataset);

auto relative_cell_pos = utils::getRelativeCellPosition(mf); // AMReX Fortran index order
std::reverse(relative_cell_pos.begin(), relative_cell_pos.end()); // now in C order
mesh_comp.setPosition(relative_cell_pos);
}
field_name += std::string("_lvl").append(std::to_string(meshLevel));
}
/*
* Write Field with all mesh levels
*
*/
void
WarpXOpenPMDPlot::WriteOpenPMDFieldsAll ( //const std::string& filename,
const std::vector<std::string>& varnames,
const amrex::Vector<amrex::MultiFab>& mf,
amrex::Vector<amrex::Geometry>& geom,
const int iteration,
const double time, bool isBTD,
const amrex::Geometry& full_BTD_snapshot ) const
{
//This is AMReX's tiny profiler. Possibly will apply it later
WARPX_PROFILE("WarpXOpenPMDPlot::WriteOpenPMDFields()");

// Loop through the multifab, and store each box as a chunk,
// in the openPMD file.
for( amrex::MFIter mfi(mf); mfi.isValid(); ++mfi ) {
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_Series != nullptr, "openPMD series must be initialized");

amrex::FArrayBox const& fab = mf[mfi];
amrex::Box const& local_box = fab.box();
// is this either a regular write (true) or the first write in a
// backtransformed diagnostic (BTD):
bool const first_write_to_iteration = ! m_Series->iterations.contains( iteration );

// Determine the offset and size of this chunk
amrex::IntVect const box_offset = local_box.smallEnd() - global_box.smallEnd();
auto const chunk_offset = getReversedVec( box_offset );
auto const chunk_size = getReversedVec( local_box.size() );
// meta data
openPMD::Iteration series_iteration = m_Series->iterations[iteration];

// Write local data
amrex::Real const * local_data = fab.dataPtr( icomp );
mesh_comp.storeChunk( openPMD::shareRaw(local_data),
chunk_offset, chunk_size );
}
auto meshes = series_iteration.meshes;
if (first_write_to_iteration) {
// lets see whether full_geom varies from geom[0] xgeom[1]
series_iteration.setTime( time );
}
// Flush data to disk after looping over all components
m_Series->flush();

for (int i=0; i<geom.size(); i++) {
amrex::Geometry full_geom = geom[i];
if( isBTD )
full_geom = full_BTD_snapshot;

// setup is called once. So it uses property "period" from first
// geometry for <all> field levels.
if ( (0 == i) && first_write_to_iteration )
SetupFields(meshes, full_geom);

amrex::Box const & global_box = full_geom.Domain();
auto const global_size = getReversedVec(global_box.size());

int const ncomp = mf[i].nComp();
for ( int icomp=0; icomp<ncomp; icomp++ ) {
std::string const & varname = varnames[icomp];

// assume fields are scalar unless they match the following match of known vector fields
std::string field_name = varname;
std::string comp_name = openPMD::MeshRecordComponent::SCALAR;
GetMeshCompNames( i, varname, field_name, comp_name );

auto mesh = meshes[field_name];
auto mesh_comp = mesh[comp_name];
if ( first_write_to_iteration )
{
SetupMeshComp( mesh, full_geom, mesh_comp );
detail::setOpenPMDUnit( mesh, field_name );

auto relative_cell_pos = utils::getRelativeCellPosition(mf[i]); // AMReX Fortran index order
std::reverse( relative_cell_pos.begin(), relative_cell_pos.end() ); // now in C order
mesh_comp.setPosition( relative_cell_pos );
}

// Loop through the multifab, and store each box as a chunk,
// in the openPMD file.
for( amrex::MFIter mfi(mf[i]); mfi.isValid(); ++mfi )
{
amrex::FArrayBox const& fab = mf[i][mfi];
amrex::Box const& local_box = fab.box();

// Determine the offset and size of this chunk
amrex::IntVect const box_offset = local_box.smallEnd() - global_box.smallEnd();
auto const chunk_offset = getReversedVec( box_offset );
auto const chunk_size = getReversedVec( local_box.size() );

amrex::Real const * local_data = fab.dataPtr( icomp );
mesh_comp.storeChunk( openPMD::shareRaw(local_data),
chunk_offset, chunk_size );
}
} // icomp loop
// Flush data to disk after looping over all components
m_Series->flush();
} // levels loop (i)
}
#endif // WARPX_USE_OPENPMD

Expand Down

0 comments on commit 45e6483

Please sign in to comment.