diff --git a/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp b/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp index eb3846e69c4..9567108c633 100644 --- a/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp +++ b/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp @@ -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); diff --git a/Source/Diagnostics/WarpXOpenPMD.H b/Source/Diagnostics/WarpXOpenPMD.H index 4152dd89692..101f51f4f30 100644 --- a/Source/Diagnostics/WarpXOpenPMD.H +++ b/Source/Diagnostics/WarpXOpenPMD.H @@ -117,10 +117,10 @@ public: void WriteOpenPMDParticles (const amrex::Vector& particle_diags); - void WriteOpenPMDFields ( + void WriteOpenPMDFieldsAll ( const std::vector& varnames, - const amrex::MultiFab& mf, - const amrex::Geometry& geom, + const amrex::Vector& mf, + amrex::Vector& geom, const int iteration, const double time, bool isBTD = false, const amrex::Geometry& full_BTD_snapshot=amrex::Geometry() ) const; @@ -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 diff --git a/Source/Diagnostics/WarpXOpenPMD.cpp b/Source/Diagnostics/WarpXOpenPMD.cpp index 2ff21de2aab..9676df501d6 100644 --- a/Source/Diagnostics/WarpXOpenPMD.cpp +++ b/Source/Diagnostics/WarpXOpenPMD.cpp @@ -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 ) ); @@ -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& 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 const grid_spacing = getReversedVec(full_geom.CellSize()); - // - Global offset - std::vector const global_offset = getReversedVec(full_geom.ProbLo()); - // - AxisLabels - std::vector axis_labels = detail::getFieldAxisLabels(); - - // Prepare the type of dataset that will be written - openPMD::Datatype const datatype = openPMD::determineDatatype(); - 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 fieldBoundary(6, "reflecting"); @@ -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 const grid_spacing = getReversedVec(full_geom.CellSize()); + // - Global offset + std::vector const global_offset = getReversedVec(full_geom.ProbLo()); + // - AxisLabels + std::vector axis_labels = detail::getFieldAxisLabels(); + + // Prepare the type of dataset that will be written + openPMD::Datatype const datatype = openPMD::determineDatatype(); + 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 + * + * @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 @@ -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& varnames, + const amrex::Vector& mf, + amrex::Vector& 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 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; icompflush(); + } // levels loop (i) } #endif // WARPX_USE_OPENPMD