added option to allow openPMD-api directly from AMReX I/O

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

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);

src/particles/elements/All.H

@@ -36,6 +36,7 @@
 #include "SoftQuad.H"
 #include "ThinDipole.H"
 #include "diagnostics/openPMD.H"
+#include "diagnostics/plotplus.H"
 
 #include <variant>
 
@@ -52,6 +53,9 @@ namespace impactx
  ChrQuad,
  ConstF,
  diagnostics::BeamMonitor,
+#ifdef AMREX_USE_OPENPMD_API
+ diagnostics::BeamPlotplus,
+#endif
  DipEdge,
  Drift,
  ExactDrift,

src/particles/elements/diagnostics/CMakeLists.txt

@@ -1,4 +1,5 @@
 target_sources(lib
  PRIVATE
  openPMD.cpp
+ plotplus.cpp
 )

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 <AMReX_PlotFileUtil.H> // for AMReXWithOpenPMD
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+
+#include <any>
+
+#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<class PIt>
+ 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<std::string>& real_names,
+ amrex::Vector<std::string>& int_names,
+ amrex::Vector<int>& int_flags,
+ amrex::Vector<int>& 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 ? "<<IntSoA::names_s.size()<<std::endl;
+ int_names.resize(IntSoA::names_s.size());
+ std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_names.begin());
+
+ for (auto int_idx=0; int_idx < IntSoA::nattribs; int_idx++) {
+ auto const component_name = int_names.at(int_idx);
+ }
+ int_flags.resize(int_names.size(), 1);
+ */
+ }
+
+ }; // class impactxWriter
+
+
+
+
+class AMReXWithOpenPMD
+{
+public:
+ AMReXWithOpenPMD();
+
+ AMReXWithOpenPMD (AMReXWithOpenPMD && other) = default;
+ AMReXWithOpenPMD& operator= (AMReXWithOpenPMD && other) = default;
+
+ ~AMReXWithOpenPMD();
+
+ void SetWriter(amrex::openpmd_api::AMReX_openPMDWriter* w);
+
+ bool InitLocalHandler(const std::string& prefix);
+
+ std::unique_ptr<amrex::openpmd_api::AMReX_openPMDHandler> 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<amrex::PinnedArenaAllocator>;
+
+ /** 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<std::string, void*> 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