Monitor: Add Reference Particle Details (#444)

* Monitor: Add Reference Particle Details

Complete all reference particle attributes for our openPMD
beam monitor. This is better than getting the details from
ASCII, because we can exactly preserve all precision. Also,
it belongs here.

In the future, we will update this to conform to the openPMD 2.0
`BeamPhysics` and `SpeciesTypes` extensions.

* Update reference particle momentum

is beta*gamma

* Doc: Update `p_t`

Co-authored-by: Chad Mitchell <[email protected]>

docs/source/dataanalysis/dataanalysis.rst

@@ -14,6 +14,40 @@ For data analysis of openPMD data, see examples of `many supported tools, Python
 
 See also `WarpX' documentation on openPMD <https://warpx.readthedocs.io/en/latest/dataanalysis/formats.html>`__.
 
+Additional Beam Attributes
+""""""""""""""""""""""""""
+
+We add the following additional attributes on the openPMD ``beam`` species at the monitor position.
+
+Reference particle:
+
+* ``beta_ref`` reference particle normalized velocity :math:`\beta = v/c`
+* ``gamma_ref`` reference particle Lorentz factor :math:`\gamma = 1/\sqrt{1-\beta^2}`
+* ``s_ref`` integrated orbit path length, in meters
+* ``x_ref`` horizontal position x, in meters
+* ``y_ref`` vertical position y, in meters
+* ``z_ref`` longitudinal position z, in meters
+* ``t_ref`` clock time * c in meters
+* ``px_ref`` momentum in x, normalized to mass*c :math:`p_x = \gamma \beta_x`
+* ``py_ref`` momentum in y, normalized to mass*c :math:`p_y = \gamma \beta_y`
+* ``pz_ref`` momentum in z, normalized to mass*c :math:`p_z = \gamma \beta_z`
+* ``pt_ref`` energy, normalized by rest energy :math:`p_t = -\gamma`
+* ``mass`` reference rest mass, in kg
+* ``charge`` reference charge, in C
+
+Example to print the integrated orbit path length ``s`` at each beam monitor position:
+
+.. code-block:: python
+
+   import openpmd_api as io
+
+   series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+
+   for k_i, i in series.iterations.items():
+       beam = i.particles["beam"]
+       s_ref = beam.get_attribute("s_ref")
+       print(f"step {k_i:>3}: s_ref={s_ref}")
+
 
 Reduced Beam Characteristics
 ----------------------------

src/particles/ReferenceParticle.H

@@ -33,10 +33,10 @@ namespace impactx
         amrex::ParticleReal y = 0.0;  ///< vertical position y, in meters
         amrex::ParticleReal z = 0.0;  ///< longitudinal position z, in meters
         amrex::ParticleReal t = 0.0;  ///< clock time * c in meters
-        amrex::ParticleReal px = 0.0; ///< momentum in x, normalized to proper velocity
-        amrex::ParticleReal py = 0.0; ///< momentum in y, normalized to proper velocity
-        amrex::ParticleReal pz = 0.0; ///< momentum in z, normalized to proper velocity
-        amrex::ParticleReal pt = 0.0; ///< energy deviation, normalized by rest energy
+        amrex::ParticleReal px = 0.0; ///< momentum in x, normalized by mass*c
+        amrex::ParticleReal py = 0.0; ///< momentum in y, normalized by mass*c
+        amrex::ParticleReal pz = 0.0; ///< momentum in z, normalized by mass*c
+        amrex::ParticleReal pt = 0.0; ///< energy, normalized by rest energy
         amrex::ParticleReal mass = 0.0; ///< reference rest mass, in kg
         amrex::ParticleReal charge = 0.0; ///< reference charge, in C
 
@@ -124,7 +124,7 @@ namespace impactx
             return *this;
         }
 
-        /** Get reference particle energy
+        /** Get reference particle kinetic energy
          *
          * @returns kinetic energy in MeV
          */

src/particles/elements/diagnostics/openPMD.H

@@ -82,6 +82,8 @@ namespace detail
         BeamMonitor& operator= (BeamMonitor && other) = default;
 
         /** Prepare entering the element before starting push logic.
+         *
+         * And write reference particle.
          *
          * @param[in] pc particle container
          * @param[in] ref_part reference particle
@@ -105,7 +107,11 @@ namespace detail
             int step
         );
 
-        // TODO: move to private function
+        /** Write a tile of particles
+         *
+         * @param pti particle tile iterator
+         * @param ref_part reference particle
+         */
         void operator() (
             PinnedContainer::ParIterType & pti,
             RefPart const & ref_part

src/particles/elements/diagnostics/openPMD.cpp

@@ -260,12 +260,22 @@ namespace detail
             }
         }
 
-        // beam mass
-        beam.setAttribute( "mass", ref_part.mass );
+        // reference particle information
         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 );
 
-        // openPMD coarse position
+        // openPMD coarse position: for global coordinates
         {
             beam["positionOffset"]["x"].resetDataset(d_fl);
             beam["positionOffset"]["x"].makeConstant(ref_part.x);
@@ -319,7 +329,7 @@ namespace detail
                           }, true);
         */
 
-        // prepare element access
+        // prepare element access & write reference particle
         this->prepare(pinned_pc, ref_part, step);
 
         // loop over refinement levels
@@ -331,10 +341,7 @@ namespace detail
             using ParIt = PinnedContainer::ParIterType;
             // note: openPMD-api is not thread-safe, so do not run OMP parallel here
             for (ParIt pti(pinned_pc, lev); pti.isValid(); ++pti) {
-                // push reference particle in global coordinates
-                this->operator()(ref_part);
-
-                // push beam particles relative to reference particle
+                // write beam particles relative to reference particle
                 this->operator()(pti, ref_part);
             } // end loop over all particle boxes
         } // end mesh-refinement level loop