diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 555efa370..232606644 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -73,6 +73,8 @@ macro(find_ablastr) # shared libs, i.e. for Python bindings, need relocatable code if(ImpactX_PYTHON OR BUILD_SHARED_LIBS) + set(AMReX_TINY_PROFILE ON CACHE BOOL "") + set(AMReX_OPENPMD_API ON CACHE BOOL "") set(AMReX_PIC ON CACHE INTERNAL "Build AMReX with position independent code") set(ABLASTR_POSITION_INDEPENDENT_CODE ON CACHE INTERNAL diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index b5f5de705..065a96d1a 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -268,6 +268,18 @@ namespace detail std::string openpmd_encoding{"g"}; pp_element.queryAdd("encoding", openpmd_encoding); m_lattice.emplace_back(diagnostics::BeamMonitor(openpmd_name, openpmd_backend, openpmd_encoding)); + } else if (element_type == "beam_plotplus") { +#ifdef AMREX_USE_OPENPMD_API + std::string openpmd_name = element_name; + pp_element.queryAdd("name", openpmd_name); + std::string openpmd_backend = "default"; + pp_element.queryAdd("backend", openpmd_backend); + std::string openpmd_encoding{"g"}; + pp_element.queryAdd("encoding", openpmd_encoding); + m_lattice.emplace_back(diagnostics::BeamPlotplus(openpmd_name, openpmd_backend, openpmd_encoding)); +#else + amrex::Abort("plotplus is not supported for lattice element " + element_name + ": " + element_type); +#endif } else if (element_type == "line") { // Parse the lattice elements for the sub-lattice in the line amrex::ParmParse pp_sub_lattice(element_name); diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index 44e6bc7e7..41db61444 100644 --- a/src/particles/elements/All.H +++ b/src/particles/elements/All.H @@ -36,6 +36,7 @@ #include "SoftQuad.H" #include "ThinDipole.H" #include "diagnostics/openPMD.H" +#include "diagnostics/plotplus.H" #include @@ -52,6 +53,9 @@ namespace impactx ChrQuad, ConstF, diagnostics::BeamMonitor, +#ifdef AMREX_USE_OPENPMD_API + diagnostics::BeamPlotplus, +#endif DipEdge, Drift, ExactDrift, diff --git a/src/particles/elements/diagnostics/CMakeLists.txt b/src/particles/elements/diagnostics/CMakeLists.txt index 757826d6a..78b50a2e0 100644 --- a/src/particles/elements/diagnostics/CMakeLists.txt +++ b/src/particles/elements/diagnostics/CMakeLists.txt @@ -1,4 +1,5 @@ target_sources(lib PRIVATE openPMD.cpp + plotplus.cpp ) diff --git a/src/particles/elements/diagnostics/plotplus.H b/src/particles/elements/diagnostics/plotplus.H new file mode 100644 index 000000000..0cf5bd9a5 --- /dev/null +++ b/src/particles/elements/diagnostics/plotplus.H @@ -0,0 +1,203 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Junmin Gu + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ELEMENTS_DIAGS_OPENPMD_WITH_AMREX_H +#define IMPACTX_ELEMENTS_DIAGS_OPENPMD_WITH_AMREX_H + +#include "particles/elements/mixin/thin.H" +#include "particles/ImpactXParticleContainer.H" + +#include // for AMReXWithOpenPMD +#include +#include + +#include + +#ifdef AMREX_USE_OPENPMD_API +namespace impactx::diagnostics +{ + class AMReX_impactxWriter: public amrex::openpmd_api::AMReX_openPMDWriter + { + public: + AMReX_impactxWriter(openPMD::IterationEncoding ie = openPMD::IterationEncoding::variableBased, + std::string filetype = "default", + std::string openpmdOptions = "{}") + :AMReX_openPMDWriter("{impactX}", ie, filetype, openpmdOptions) + {} + + + // subclasses + inline std::vector< std::string > getParticlePositionComponentLabels() const + { + using vs = std::vector< std::string >; + vs const positionComponents{"x", "y", "t"}; + return positionComponents; + } + // internal to impact x + /** Prepare entering the element before starting push logic. + * + * @param[in] pc particle container + * @param[in] ref_part reference particle + * @param[in] step global step for diagnostics + */ + void SetConstantRefPart (openPMD::ParticleSpecies& beam, + //const unsigned long long& np, + RefPart const & ref_part + ) + { + // beam mass + beam.setAttribute( "beta_ref", ref_part.beta() ); + beam.setAttribute( "gamma_ref", ref_part.gamma() ); + beam.setAttribute( "s_ref", ref_part.s ); + beam.setAttribute( "x_ref", ref_part.x ); + beam.setAttribute( "y_ref", ref_part.y ); + beam.setAttribute( "z_ref", ref_part.z ); + beam.setAttribute( "t_ref", ref_part.t ); + beam.setAttribute( "px_ref", ref_part.px ); + beam.setAttribute( "py_ref", ref_part.py ); + beam.setAttribute( "pz_ref", ref_part.pz ); + beam.setAttribute( "pt_ref", ref_part.pt ); + beam.setAttribute( "mass", ref_part.mass ); + beam.setAttribute( "charge", ref_part.charge ); + + beam["positionOffset"]["x"].makeConstant( ref_part.x ); + beam["positionOffset"]["y"].makeConstant( ref_part.y ); + beam["positionOffset"]["t"].makeConstant( ref_part.t ); + + } + + template + void Save_impactx_PosID(PIt& pti, + openPMD::ParticleSpecies& currSpecies, + unsigned long long offset) const + { + // use default is enough I guess? + SavePosId(pti, currSpecies, offset); + } + + void + GetNames(amrex::Vector& real_names, + amrex::Vector& int_names, + amrex::Vector& int_flags, + amrex::Vector& real_flags) const + + { + real_names.resize(RealSoA::names_s.size()); + std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_names.begin()); + + for (auto real_idx=0; real_idx < RealSoA::nattribs; real_idx++) { + auto const component_name = real_names.at(real_idx); + } + real_flags.resize(real_names.size(), 1); + + + // currently no use of int names + int_names.resize(0); + int_flags.resize(0); + + /* + // enable when IntSoA is in + amrex::Print() <<" => [check] int names ? "< m_UserHandler; + +private: + std::string m_Prefix = ""; +}; + + + + + /** This element writes the particle beam out to openPMD data. + * + * This class behaves like a singleton if constructed with the + * same series name as an existing instance. + */ + struct BeamPlotplus + : public elements::Thin + { + static constexpr auto name = "BeamPlotplus"; + using PinnedContainer = typename ImpactXParticleContainer::ContainerLike; + + /** This element writes the particle beam out to openPMD data. + * + * Elements with the same series name are identical. + * + * @param series_name name of the data series, usually the element name + * @param backend file format backend for openPMD, e.g., "bp" or "h5" + * @param encoding openPMD iteration encoding: "v"ariable based, "f"ile based, "g"roup based (default) + */ + BeamPlotplus (std::string series_name, std::string backend="default", std::string encoding="g"); + + BeamPlotplus (BeamPlotplus const & other) = default; + BeamPlotplus (BeamPlotplus && other) = default; + BeamPlotplus& operator= (BeamPlotplus const & other) = default; + BeamPlotplus& operator= (BeamPlotplus && other) = default; + + + /** Dump all particles. + * + * Particles are relative to the reference particle. + * + * @param[in,out] pc particle container to push + * @param[in] step global step for diagnostics + */ + void operator() ( + ImpactXParticleContainer & pc, + int step + ); + + /** This does nothing to the reference particle. */ + using Thin::operator(); + + /** Close and deallocate all data series and backends. + */ + void + finalize (); + + private: + std::string m_seriesName; + void* m_plotWriter = NULL; + static inline std::map m_uniqueWriter = {}; + + //int m_file_min_digits = 6; //! minimum number of digits to iteration number in file name + }; + +} // namespace impactx::diagnostics + +#endif // AMREX_USE_OPENPMD_API +#endif // IMPACTX_ELEMENTS_DIAGS_OPENPMD_WITH_AMREX_H diff --git a/src/particles/elements/diagnostics/plotplus.cpp b/src/particles/elements/diagnostics/plotplus.cpp new file mode 100644 index 000000000..9d06a3605 --- /dev/null +++ b/src/particles/elements/diagnostics/plotplus.cpp @@ -0,0 +1,220 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ + +/// #include "openPMD.H" +#include "ImpactXVersion.H" +#include "particles/ImpactXParticleContainer.H" + +#include + +#include +#include +#include + +#include + +#ifdef AMREX_USE_OPENPMD_API +#include "plotplus.H" + +namespace impactx::diagnostics +{ + + class StepMgr +{ +public: + StepMgr(int step, AMReXWithOpenPMD* owner) + :m_Step(step), + m_Owner(owner) + { + m_Owner = owner; + m_Owner->m_UserHandler->m_Writer->SetStep(m_Step); + } + ~StepMgr() + { + m_Owner->m_UserHandler->m_Writer->CloseStep(m_Step); + } + +private: + int m_Step; + AMReXWithOpenPMD* m_Owner; +}; + + + AMReXWithOpenPMD::AMReXWithOpenPMD() + { + // warpx has multiple diags, each should maintain its own handler + m_UserHandler = amrex::openpmd_api::InitUserHandler(m_Prefix); + } + + void AMReXWithOpenPMD::SetWriter(amrex::openpmd_api::AMReX_openPMDWriter* w) + { + BL_ASSERT ( m_UserHandler != nullptr ); + BL_ASSERT ( w != nullptr ); + + m_UserHandler->SetWriter(w); + } + + + AMReXWithOpenPMD::~AMReXWithOpenPMD() +{ + amrex::openpmd_api::CloseUserHandler(m_UserHandler); +} + +bool AMReXWithOpenPMD::InitLocalHandler(const std::string& prefix) +{ + if (m_Prefix.compare(prefix) == 0) + return false; + + amrex::Print()<<" openpmd_api::Init handler "< 0u) + { + m_plotWriter = m_uniqueWriter[m_seriesName]; + } + else + { + auto m_Writer = new AMReXWithOpenPMD(); +#ifdef ImpactX_USE_OPENPMD + // encoding of iterations in the series + openPMD::IterationEncoding series_encoding = openPMD::IterationEncoding::groupBased; + if ( "v" == encoding ) + series_encoding = openPMD::IterationEncoding::variableBased; + else if ( "g" == encoding ) + series_encoding = openPMD::IterationEncoding::groupBased; + else if ( "f" == encoding ) + series_encoding = openPMD::IterationEncoding::fileBased; + + if ( m_Writer->InitLocalHandler(m_seriesName) ) + { + AMReX_impactxWriter* testWriter = new AMReX_impactxWriter(series_encoding); + m_Writer->SetWriter(testWriter); + } +#else + amrex::AllPrint() << "Warning: openPMD output requested but not compiled for series=" << m_series_name << "\n"; +#endif + m_plotWriter = m_Writer; + m_uniqueWriter[m_seriesName] = m_plotWriter; + }// else + } + + #ifdef NEVER + void BeamPlotplus::prepare ( + PinnedContainer & pc, + RefPart const & ref_part, + int step + ) { +#ifdef ImpactX_USE_OPENPMD + + /* should be covered by amrex-openpmd-io + // SoA: Real + { + std::vector real_soa_names(RealSoA::names_s.size()); + std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); + for (auto real_idx = 0; real_idx < RealSoA::nattribs; real_idx++) { + auto const component_name = real_soa_names.at(real_idx); + getComponentRecord(component_name).resetDataset(d_fl); + } + } + // SoA: Int + static_assert(IntSoA::nattribs == 0); // not yet used + */ +#else + amrex::ignore_unused(pc, step); +#endif + } +#endif + void + BeamPlotplus::operator() ( + ImpactXParticleContainer & pc, + int step + ) + { + BL_PROFILE("BeamPlotplus::(pc, step)") + // preparing to access reference particle data: RefPart + RefPart & ref_part = pc.GetRefParticle(); + + // pinned memory copy + PinnedContainer pinned_pc = pc.make_alike(); + pinned_pc.copyParticles(pc, true); // no filtering + + // TODO: filtering + /* + using SrcData = WarpXParticleContainer::ParticleTileType::ConstParticleTileDataType; + tmp.copyParticles(*pc, + [=] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip, const amrex::RandomEngine& engine) + { + const SuperParticleType& p = src.getSuperParticle(ip); + return random_filter(p, engine) * uniform_filter(p, engine) + * parser_filter(p, engine) * geometry_filter(p, engine); + }, true); + */ + + auto m_Writer = (AMReXWithOpenPMD*)(m_plotWriter); + + StepMgr sm(step, m_Writer); + pinned_pc.CountParticles(); + + + AMReX_impactxWriter* impactxWriter = (AMReX_impactxWriter*) (m_Writer->m_UserHandler->m_Writer.get()); + amrex::Vector real_names; + amrex::Vector int_names; + amrex::Vector int_flags; + amrex::Vector real_flags; + impactxWriter->GetNames(real_names, int_names, int_flags, real_flags); + m_Writer->m_UserHandler->m_Writer->DumpParticles(pinned_pc, + "beam", + real_flags, + int_flags, + real_names, + int_names, + + [=] ([[maybe_unused]] auto& ppc, openPMD::ParticleSpecies& currSpecies, [[maybe_unused]] unsigned long long localTotal) + { + impactxWriter->SetConstantRefPart(currSpecies, ref_part); + }, + [=] (auto& pti, openPMD::ParticleSpecies& currSpecies, unsigned long long offset) + { + // use the default + impactxWriter->Save_impactx_PosID(pti, currSpecies, offset); + }); + + + } +} // namespace impactx::diagnostics + + +#endif //#ifdef AMREX_USE_OPENPMD_API