Improve the Handling of Runtime Arguments

Simplify and remove hacks as long as AMReX has no upstream
support for named SoA components.

src/ImpactX.cpp

@@ -90,7 +90,7 @@ namespace impactx {
         // alloc particle containers
         //   the lost particles have an extra runtime attribute: s when it was lost
         bool comm = true;
-        amr_data->m_particles_lost->AddRealComp(comm);
+        amr_data->m_particles_lost->AddRealComp("s_lost", comm);
 
         //   have to resize here, not in the constructor because grids have not
         //   been built when constructor was called.

src/particles/ImpactXParticleContainer.H

@@ -100,6 +100,8 @@ namespace impactx
         ParIterSoA (ContainerType& pc, int level);
 
         ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
+
+        ContainerType * pc () { return m_pc; }
     };
 
     /** Const AMReX iterator for particle boxes - data is read only.
@@ -116,6 +118,8 @@ namespace impactx
         ParConstIterSoA (ContainerType& pc, int level);
 
         ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
+
+        ContainerType const * pc () const { return m_pc; }
     };
 
     /** Beam Particles in ImpactX
@@ -138,6 +142,20 @@ namespace impactx
         //! Destruct a particle container
         virtual ~ImpactXParticleContainer() = default;
 
+        /** Add a ParticleReal component
+         *
+         * @param name a unique name of the component, not colliding with RealSoA
+         * @param communicate participate in MPI communication when particles move
+         */
+        void AddRealComp (std::string const & name, bool communicate=true);
+
+        /** Add an Integer component
+         *
+         * @param name a unique name of the component, not colliding with IntSoA
+         * @param communicate participate in MPI communication when particles move
+         */
+        void AddIntComp (std::string const & name, bool communicate=true);
+
         /** Add new particles to the container for fixed s.
          *
          * Note: This can only be used *after* the initialization (grids) have
@@ -291,23 +309,13 @@ namespace impactx
         //! the current coordinate system of particles in this container
         CoordSystem m_coordsystem = CoordSystem::s;
 
-    }; // ImpactXParticleContainer
+        //! ParticleReal component names
+        std::vector<std::string> m_real_soa_names;
 
-    /** Get the name of each ParticleReal SoA component
-     *
-     * @param num_real_comps number of compile-time + runtime arrays
-     * @return names
-     */
-    std::vector<std::string>
-    get_RealSoA_names (int num_real_comps);
+        //! Int component names
+        std::vector<std::string> m_int_soa_names;
 
-    /** Get the name of each int SoA component
-     *
-     * @param num_int_comps number of compile-time + runtime arrays
-     * @return names
-     */
-    std::vector<std::string>
-    get_intSoA_names (int num_int_comps);
+    }; // ImpactXParticleContainer
 
 } // namespace impactx
 

src/particles/ImpactXParticleContainer.cpp

@@ -58,6 +58,26 @@ namespace impactx
         : amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
     {
         SetParticleSize();
+
+        // name compile-time attributes
+        m_real_soa_names.resize(RealSoA::names_s.size());
+        m_int_soa_names.resize(IntSoA::names_s.size());
+        std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), m_real_soa_names.begin());
+        std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), m_int_soa_names.begin());
+    }
+
+    void
+    ImpactXParticleContainer::AddRealComp (std::string const & name, bool communicate)
+    {
+        m_real_soa_names.push_back(name);
+        amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>::AddRealComp(communicate);
+    }
+
+    void
+    ImpactXParticleContainer::AddIntComp (std::string const & name, bool communicate)
+    {
+        m_int_soa_names.push_back(name);
+        amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>::AddIntComp(communicate);
     }
 
     void
@@ -246,13 +266,13 @@ namespace impactx
     std::vector<std::string>
     ImpactXParticleContainer::RealSoA_names () const
     {
-        return get_RealSoA_names(this->NumRealComps());
+        return m_real_soa_names;
     }
 
     std::vector<std::string>
     ImpactXParticleContainer::intSoA_names () const
     {
-        return get_intSoA_names(this->NumIntComps());
+        return m_int_soa_names;
     }
 
     CoordSystem
@@ -266,40 +286,4 @@ namespace impactx
     {
         m_coordsystem = coord_system;
     }
-
-    std::vector<std::string>
-    get_RealSoA_names (int num_real_comps)
-    {
-        std::vector<std::string> real_soa_names(num_real_comps);
-
-        // compile-time attributes
-        std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin());
-
-        // runtime attributes
-        if (num_real_comps > int(RealSoA::names_s.size()))
-        {
-            // particles lost record their "s" position where they got lost
-            real_soa_names[RealSoA::nattribs] = "s_lost";
-        }
-
-        return real_soa_names;
-    }
-
-    std::vector<std::string>
-    get_intSoA_names (int num_int_comps)
-    {
-        std::vector<std::string> int_soa_names(num_int_comps);
-
-        // compile-time attributes
-        std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_soa_names.begin());
-
-        // runtime attributes
-        if (num_int_comps > int(IntSoA::names_s.size()))
-        {
-            // particles lost record their "s" position where they got lost
-            int_soa_names[IntSoA::nattribs] = "s_lost";
-        }
-
-        return int_soa_names;
-    }
 } // namespace impactx

src/particles/elements/diagnostics/openPMD.H

@@ -17,6 +17,8 @@
 #include <AMReX_REAL.H>
 
 #include <any>
+#include <string>
+#include <vector>
 
 
 namespace impactx::diagnostics
@@ -86,11 +88,15 @@ namespace detail
          * And write reference particle.
          *
          * @param[in] pc particle container
+         * @param[in] real_soa_names ParticleReal component names
+         * @param[in] int_soa_names integer component names
          * @param[in] ref_part reference particle
          * @param[in] step global step for diagnostics
          */
         void prepare (
             PinnedContainer & pc,
+            std::vector<std::string> const & real_soa_names,
+            std::vector<std::string> const & int_soa_names,
             RefPart const & ref_part,
             int step
         );
@@ -110,10 +116,14 @@ namespace detail
         /** Write a tile of particles
          *
          * @param pti particle tile iterator
+         * @param[in] real_soa_names ParticleReal component names
+         * @param[in] int_soa_names integer component names
          * @param ref_part reference particle
          */
         void operator() (
             PinnedContainer::ParIterType & pti,
+            std::vector<std::string> const & real_soa_names,
+            std::vector<std::string> const & int_soa_names,
             RefPart const & ref_part
         );
 

src/particles/elements/diagnostics/openPMD.cpp

@@ -24,6 +24,7 @@ namespace io = openPMD;
 
 #include <string>
 #include <utility>
+#include <vector>
 
 
 namespace impactx::diagnostics
@@ -217,6 +218,8 @@ namespace detail
 
     void BeamMonitor::prepare (
         PinnedContainer & pc,
+        std::vector<std::string> const & real_soa_names,
+        std::vector<std::string> const & int_soa_names,
         RefPart const & ref_part,
         int step
     ) {
@@ -281,14 +284,15 @@ namespace detail
 
         // SoA: Real
         {
-            std::vector<std::string> real_soa_names = get_RealSoA_names(pc.NumRealComps());
             for (auto real_idx = 0; real_idx < pc.NumRealComps(); 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
+        if (!int_soa_names.empty())
+            throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!");
 #else
         amrex::ignore_unused(pc, step);
 #endif
@@ -303,6 +307,10 @@ namespace detail
         std::string profile_name = "impactx::Push::" + std::string(BeamMonitor::name);
         BL_PROFILE(profile_name);
 
+        // component names
+        std::vector<std::string> real_soa_names = pc.RealSoA_names();
+        std::vector<std::string> int_soa_names = pc.intSoA_names();
+
         // preparing to access reference particle data: RefPart
         RefPart & ref_part = pc.GetRefParticle();
 
@@ -323,7 +331,7 @@ namespace detail
         */
 
         // prepare element access & write reference particle
-        this->prepare(pinned_pc, ref_part, step);
+        this->prepare(pinned_pc, real_soa_names, int_soa_names, ref_part, step);
 
         // loop over refinement levels
         int const nLevel = pinned_pc.finestLevel();
@@ -335,7 +343,7 @@ namespace detail
             // note: openPMD-api is not thread-safe, so do not run OMP parallel here
             for (ParIt pti(pinned_pc, lev); pti.isValid(); ++pti) {
                 // write beam particles relative to reference particle
-                this->operator()(pti, ref_part);
+                this->operator()(pti, real_soa_names, int_soa_names, ref_part);
             } // end loop over all particle boxes
         } // end mesh-refinement level loop
 
@@ -350,6 +358,8 @@ namespace detail
     void
     BeamMonitor::operator() (
         PinnedContainer::ParIterType & pti,
+        std::vector<std::string> const & real_soa_names,
+        std::vector<std::string> const & int_soa_names,
         RefPart const & ref_part
     )
     {
@@ -386,7 +396,6 @@ namespace detail
         }
         //   SoA floating point (ParticleReal) properties
         {
-            std::vector<std::string> real_soa_names = get_RealSoA_names(soa.NumRealComps());
             for (auto real_idx=0; real_idx < soa.NumRealComps(); real_idx++) {
                 auto const component_name = real_soa_names.at(real_idx);
                 getComponentRecord(component_name).storeChunkRaw(
@@ -396,10 +405,11 @@ namespace detail
         //   SoA integer (int) properties (not yet used)
         {
             static_assert(IntSoA::nattribs == 0); // not yet used
+            if (!int_soa_names.empty())
+                throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!");
             /*
             // comment this in once IntSoA::nattribs is > 0
 
-            std::vector<std::string> int_soa_names(IntSoA::names_s.size);
             std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_soa_names.begin());
 
             for (auto int_idx=0; int_idx < RealSoA::nattribs; int_idx++) {

src/python/ImpactXParticleContainer.cpp

@@ -30,7 +30,8 @@ void init_impactxparticlecontainer(py::module& m)
     py::class_<
         ParIterSoA,
         amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>
-    >(m, "ImpactXParIter")
+    > py_pariter_soa(m, "ImpactXParIter");
+    py_pariter_soa
         .def(py::init<ParIterSoA::ContainerType&, int>(),
              py::arg("particle_container"), py::arg("level"))
         .def(py::init<ParIterSoA::ContainerType&, int, amrex::MFItInfo&>(),
@@ -40,7 +41,8 @@ void init_impactxparticlecontainer(py::module& m)
     py::class_<
         ParConstIterSoA,
         amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>
-    >(m, "ImpactXParConstIter")
+    > py_parconstiter_soa(m, "ImpactXParConstIter");
+    py_parconstiter_soa
         .def(py::init<ParConstIterSoA::ContainerType&, int>(),
              py::arg("particle_container"), py::arg("level"))
         .def(py::init<ParConstIterSoA::ContainerType&, int, amrex::MFItInfo&>(),
@@ -140,10 +142,6 @@ void init_impactxparticlecontainer(py::module& m)
                "Get the name of each int SoA component")
     ;
 
-    m.def("get_RealSoA_names", &get_RealSoA_names,
-          py::arg("num_real_comps"),
-          "Get the name of each ParticleReal SoA component\n\nnum_real_comps: pass number of compile-time + runtime arrays");
-    m.def("get_intSoA_names", &get_intSoA_names,
-          py::arg("num_int_comps"),
-          "Get the name of each int SoA component\n\nnum_int_comps: pass number of compile-time + runtime arrays");
+    py_pariter_soa.def("pc", &ParIterSoA::pc);
+    py_parconstiter_soa.def("pc", &ParConstIterSoA::pc);
 }