openPMD: Skip 0-positions for deselected comps

Author: ax3l

Source/Diagnostics/WarpXOpenPMD.H

@@ -218,11 +218,13 @@ private:
   /** This function sets up the entries for storing the particle positions and global IDs
   *
   * @param[in] currSpecies Corresponding openPMD species
+  * @param[in] positionComponents user-selected components of the particle position
   * @param[in] np          Number of particles
   * @param[in] isBTD       Is this a back-transformed diagnostics output?
   */
   void SetupPos (
         openPMD::ParticleSpecies& currSpecies,
+        std::vector<std::string> const & positionComponents,
         const unsigned long long& np,
         bool isBTD = false);
 
@@ -231,12 +233,14 @@ private:
    * Sets the entries for storing particle position offset, constant records (charge, mass) and ED-PIC attributes.
    *
    * @param[in] currSpecies Corresponding openPMD species
+   * @param[in] positionComponents user-selected components of the particle position
    * @param[in] np          Number of particles
    * @param[in] charge      Charge of the particles (note: fix for ions)
    * @param[in] mass        Mass of the particles
    */
   void SetConstParticleRecordsEDPIC (
         openPMD::ParticleSpecies& currSpecies,
+        std::vector<std::string> const & positionComponents,
         const unsigned long long& np,
         amrex::ParticleReal charge,
         amrex::ParticleReal mass);
@@ -284,8 +288,8 @@ private:
    * @param[in] name species name
    * @param[in] iteration timestep
    * @param[in] write_real_comp The real attribute ids, from WarpX
-   * @param[in] real_comp_names The real attribute names, from WarpX
    * @param[in] write_int_comp The int attribute ids, from WarpX
+   * @param[in] real_comp_names The real attribute names, from WarpX
    * @param[in] int_comp_names The int attribute names, from WarpX
    * @param[in] charge         Charge of the particles (note: fix for ions)
    * @param[in] mass           Mass of the particles

Source/Diagnostics/WarpXOpenPMD.cpp

@@ -213,32 +213,29 @@ namespace detail
         return make_pair(record_name, component_name);
     }
 
-    /** Return the component labels for particle positions
+    /** Return the user-selected components for particle positions
+     *
+     * @param[in] write_real_comp The real attribute ids, from WarpX
+     * @param[in] real_comp_names The real attribute names, from WarpX
      */
     inline std::vector< std::string >
-    getParticlePositionComponentLabels (bool ignore_dims=false)
+    getParticlePositionComponentLabels (
+        amrex::Vector<int> const & write_real_comp,
+        amrex::Vector<std::string> const & real_comp_names)
     {
-        using vs = std::vector< std::string >;
-        auto positionComponents = vs{"x", "y", "z"};
-        if (!ignore_dims) {
-#if defined(WARPX_DIM_1D_Z)
-            positionComponents = vs{"z"};
-#elif defined(WARPX_DIM_XZ)
-            positionComponents = vs{"x", "z"};
-#elif defined(WARPX_DIM_RZ)
-            // note: although we internally store particle positions
-            //       for AMReX in r,z and a theta attribute, we
-            //       actually need them for algorithms (e.g. push)
-            //       and I/O in Cartesian.
-            //       Other attributes like momentum are consequently
-            //       stored in x,y,z internally.
-            positionComponents = vs{"x", "y", "z"};
-#elif defined(WARPX_DIM_3D)
-            positionComponents = vs{"x", "y", "z"};
-#else
-#   error Unknown WarpX dimensionality.
-#endif
+        std::vector< std::string > positionComponents;
+
+        int idx = 0;
+        for (auto const & comp : real_comp_names ) {
+            if (write_real_comp[idx]) {
+                if (comp == "position_x" || comp == "position_y" || comp == "position_z") {
+                    std::string const last_letter{comp.back()};
+                    positionComponents.push_back(last_letter);
+                }
+            }
+            idx++;
         }
+
         return positionComponents;
     }
 
@@ -705,16 +702,18 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
         doParticleSetup = is_first_flush_with_particles || is_last_flush_and_never_particles;
     }
 
+    auto const positionComponents = detail::getParticlePositionComponentLabels(write_real_comp, real_comp_names);
+
     // this setup stage also implicitly calls "makeEmpty" if needed (i.e., is_last_flush_and_never_particles)
     //   for BTD, we call this multiple times as we may resize in subsequent dumps if number of particles in the buffer > 0
     if (doParticleSetup || is_resizing_flush) {
-        SetupPos(currSpecies, NewParticleVectorSize, isBTD);
+        SetupPos(currSpecies, positionComponents, NewParticleVectorSize, isBTD);
         SetupRealProperties(pc, currSpecies, write_real_comp, real_comp_names, write_int_comp, int_comp_names,
                             NewParticleVectorSize, isBTD);
     }
 
     if (is_last_flush_to_step) {
-        SetConstParticleRecordsEDPIC(currSpecies, NewParticleVectorSize, charge, mass);
+        SetConstParticleRecordsEDPIC(currSpecies, positionComponents, NewParticleVectorSize, charge, mass);
     }
 
     // open files from all processors, in case some will not contribute below
@@ -978,6 +977,7 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
 void
 WarpXOpenPMDPlot::SetupPos (
     openPMD::ParticleSpecies& currSpecies,
+    std::vector<std::string> const & positionComponents,
     const unsigned long long& np,
     bool const isBTD)
 {
@@ -986,7 +986,6 @@ WarpXOpenPMDPlot::SetupPos (
     auto realType = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np}, options);
     auto idType = openPMD::Dataset(openPMD::determineDatatype< uint64_t >(), {np}, options);
 
-    auto const positionComponents = detail::getParticlePositionComponentLabels();
     for( auto const& comp : positionComponents ) {
         currSpecies["position"][comp].resetDataset( realType );
     }
@@ -998,6 +997,7 @@ WarpXOpenPMDPlot::SetupPos (
 void
 WarpXOpenPMDPlot::SetConstParticleRecordsEDPIC (
         openPMD::ParticleSpecies& currSpecies,
+        std::vector<std::string> const & positionComponents,
         const unsigned long long& np,
         amrex::ParticleReal const charge,
         amrex::ParticleReal const mass)
@@ -1006,7 +1006,6 @@ WarpXOpenPMDPlot::SetConstParticleRecordsEDPIC (
     const auto *const scalar = openPMD::RecordComponent::SCALAR;
 
     // define record shape to be number of particles
-    auto const positionComponents = detail::getParticlePositionComponentLabels(true);
     for( auto const& comp : positionComponents ) {
         currSpecies["positionOffset"][comp].resetDataset( realType );
     }
@@ -1037,16 +1036,20 @@ WarpXOpenPMDPlot::SetConstParticleRecordsEDPIC (
 #endif
 
     // meta data
-    currSpecies["position"].setUnitDimension( detail::getUnitDimension("position") );
-    currSpecies["positionOffset"].setUnitDimension( detail::getUnitDimension("positionOffset") );
+    if (!positionComponents.empty()) {
+        currSpecies["position"].setUnitDimension( detail::getUnitDimension("position") );
+        currSpecies["positionOffset"].setUnitDimension( detail::getUnitDimension("positionOffset") );
+    }
     currSpecies["charge"].setUnitDimension( detail::getUnitDimension("charge") );
     currSpecies["mass"].setUnitDimension( detail::getUnitDimension("mass") );
 
     // meta data for ED-PIC extension
-    currSpecies["position"].setAttribute( "macroWeighted", 0u );
-    currSpecies["position"].setAttribute( "weightingPower", 0.0 );
-    currSpecies["positionOffset"].setAttribute( "macroWeighted", 0u );
-    currSpecies["positionOffset"].setAttribute( "weightingPower", 0.0 );
+    if (!positionComponents.empty()) {
+        currSpecies["position"].setAttribute( "macroWeighted", 0u );
+        currSpecies["position"].setAttribute( "weightingPower", 0.0 );
+        currSpecies["positionOffset"].setAttribute( "macroWeighted", 0u );
+        currSpecies["positionOffset"].setAttribute( "weightingPower", 0.0 );
+    }
     currSpecies["id"].setAttribute( "macroWeighted", 0u );
     currSpecies["id"].setAttribute( "weightingPower", 0.0 );
     currSpecies["charge"].setAttribute( "macroWeighted", 0u );