Particle Container: Track s or t (#497)

Track the current coordinate system `s` or `t` as meta data in the particle container. Throw an exception if an invalid transformation is attempted. This prevents that we accidentially transform "twice" into a direction, as happend to us in the past in interactive/scripting use from Python.
ECP-WarpX · Jan 6, 2024 · b8169b4 · b8169b4
1 parent 71df3bd
commit b8169b4
Show file tree

Hide file tree

Showing 9 changed files with 77 additions and 25 deletions.
diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp
@@ -186,7 +186,7 @@ namespace impactx
                         // transform from x',y',t to x,y,z
                         transformation::CoordinateTransformation(
                                 *m_particle_container,
-                                transformation::Direction::to_fixed_t);
+                                CoordSystem::t);
 
                         // Note: The following operation assume that
                         // the particles are in x, y, z coordinates.
@@ -218,7 +218,7 @@ namespace impactx
 
                         // transform from x,y,z to x',y',t
                         transformation::CoordinateTransformation(*m_particle_container,
-                                                                 transformation::Direction::to_fixed_s);
+                                                                 CoordSystem::s);
                     }
 
                     // for later: original Impact implementation as an option

diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H
@@ -28,6 +28,12 @@
 
 namespace impactx
 {
+    enum class CoordSystem
+    {
+        s,  ///< fixed s as the independent variable
+        t   ///< fixed t as the independent variable
+    };
+
     /** AMReX pre-defined Real attributes
      *
      * These are the AMReX pre-defined struct indexes for the Real attributes
@@ -277,6 +283,17 @@ namespace impactx
         std::vector<std::string>
         RealSoA_names () const;
 
+        /** Get the current coordinate system of particles in this container */
+        CoordSystem
+        GetCoordSystem () const;
+
+        /** Set the current coordinate system of particles in this container
+         *
+         * @param coord_system the new coordinate system the beam particles are in
+         */
+        void
+        SetCoordSystem (CoordSystem coord_system);
+
       private:
 
         //! the reference particle for the beam in the particle container
@@ -288,6 +305,9 @@ namespace impactx
         //! a non-owning reference to lost particles, i.e., due to apertures
         ImpactXParticleContainer* m_particles_lost = nullptr;
 
+        //! the current coordinate system of particles in this container
+        CoordSystem m_coordsystem = CoordSystem::s;
+
     }; // ImpactXParticleContainer
 
     /** Get the name of each Real AoS component */

diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp
@@ -237,6 +237,18 @@ namespace impactx
         return get_RealSoA_names(this->NumRealComps());
     }
 
+    CoordSystem
+    ImpactXParticleContainer::GetCoordSystem () const
+    {
+        return m_coordsystem;
+    }
+
+    void
+    ImpactXParticleContainer::SetCoordSystem (CoordSystem coord_system)
+    {
+        m_coordsystem = coord_system;
+    }
+
     std::vector<std::string>
     get_RealAoS_names ()
     {

diff --git a/src/particles/transformation/CoordinateTransformation.H b/src/particles/transformation/CoordinateTransformation.H
@@ -15,21 +15,13 @@
 
 namespace impactx::transformation
 {
-    /** Direction of the \see CoordinateTransformation
-     */
-    enum class Direction
-    {
-        to_fixed_s,
-        to_fixed_t
-    };
-
     /** Transform the coordinate system of all particles
      *
      * @param pc container of the particles to push
      * @param direction the transformation (to fixed s or to fixed t)
      */
-    void CoordinateTransformation (ImpactXParticleContainer &pc,
-                                   Direction const & direction);
+    void CoordinateTransformation (ImpactXParticleContainer & pc,
+                                   CoordSystem direction);
 
 } // namespace impactx::transformation
 

diff --git a/src/particles/transformation/CoordinateTransformation.cpp b/src/particles/transformation/CoordinateTransformation.cpp
@@ -21,12 +21,18 @@
 
 namespace impactx::transformation
 {
-    void CoordinateTransformation (ImpactXParticleContainer &pc,
-                                   Direction const &direction)
+    void CoordinateTransformation (ImpactXParticleContainer & pc,
+                                   CoordSystem direction)
     {
         BL_PROFILE("impactx::transformation::CoordinateTransformation");
         using namespace amrex::literals; // for _rt and _prt
 
+        if (direction == CoordSystem::s) {
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(pc.GetCoordSystem() == CoordSystem::t, "Already in fixed s coordinates!");
+        } else {
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(pc.GetCoordSystem() == CoordSystem::s, "Already in fixed t coordinates!");
+        }
+
         // preparing to access reference particle data: RefPart
         RefPart const ref_part = pc.GetRefParticle();
         amrex::ParticleReal const pd = ref_part.pt;  // Design value of pt/mc2 = -gamma
@@ -52,7 +58,7 @@ namespace impactx::transformation
                 amrex::ParticleReal *const AMREX_RESTRICT part_px = soa_real[RealSoA::px].dataPtr();
                 amrex::ParticleReal *const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr();
 
-                if( direction == Direction::to_fixed_s) {
+                if (direction == CoordSystem::s) {
                     BL_PROFILE("impactx::transformation::CoordinateTransformation::to_fixed_s");
 
                     amrex::ParticleReal *const AMREX_RESTRICT part_pz = soa_real[RealSoA::pz].dataPtr();
@@ -93,5 +99,8 @@ namespace impactx::transformation
                 }
             } // end loop over all particle boxes
         } // env mesh-refinement level loop
+
+        // update coordinate system meta data
+        pc.SetCoordSystem(direction);
     }
 } // namespace impactx::transformation
diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp
@@ -22,6 +22,11 @@ using namespace impactx;
 
 void init_impactxparticlecontainer(py::module& m)
 {
+    py::enum_<CoordSystem>(m, "CoordSystem")
+        .value("s", CoordSystem::s)
+        .value("t", CoordSystem::t)
+        .export_values();
+
     py::class_<
         ParIter,
         amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>
@@ -57,6 +62,11 @@ void init_impactxparticlecontainer(py::module& m)
             "RealSoA attribute name labels"
         )
 
+        .def_property_readonly("coord_system",
+            &ImpactXParticleContainer::GetCoordSystem,
+            "Get the current coordinate system of particles in this container"
+        )
+
         .def("add_n_particles",
              &ImpactXParticleContainer::AddNParticles,
              py::arg("x"), py::arg("y"), py::arg("t"),

diff --git a/src/python/transformation.cpp b/src/python/transformation.cpp
@@ -15,11 +15,8 @@ void init_transformation(py::module& m)
 {
     m.def("coordinate_transformation",
         &transformation::CoordinateTransformation,
+        py::arg("pc"),
+        py::arg("direction"),
         "Transform coordinates from fixed s to fixed to or vice versa."
     );
-
-    py::enum_<transformation::Direction>(m, "TransformationDirection")
-        .value("to_fixed_s", transformation::Direction::to_fixed_s)
-        .value("to_fixed_t", transformation::Direction::to_fixed_t)
-        .export_values();
 }
diff --git a/tests/python/test_transformation.py b/tests/python/test_transformation.py
@@ -7,13 +7,14 @@
 # -*- coding: utf-8 -*-
 
 import numpy as np
+import pytest
 
 from impactx import (
     Config,
+    CoordSystem,
     ImpactX,
     ImpactXParIter,
     RefPart,
-    TransformationDirection,
     coordinate_transformation,
     distribution,
     elements,
@@ -61,11 +62,22 @@ def test_transformation():
         mutpt=0.8,
     )
     sim.add_particles(bunch_charge_C, distr, npart)
-
     rbc_s0 = pc.reduced_beam_characteristics()
-    coordinate_transformation(pc, TransformationDirection.to_fixed_t)
+
+    # this must fail: we cannot transform from s to s
+    with pytest.raises(Exception):
+        coordinate_transformation(pc, direction=CoordSystem.s)
+
+    # transform to t
+    coordinate_transformation(pc, direction=CoordSystem.t)
     rbc_t = pc.reduced_beam_characteristics()
-    coordinate_transformation(pc, TransformationDirection.to_fixed_s)
+
+    # this must fail: we cannot transform from t to t
+    with pytest.raises(Exception):
+        coordinate_transformation(pc, direction=CoordSystem.t)
+
+    # transform back to s
+    coordinate_transformation(pc, direction=CoordSystem.s)
     rbc_s = pc.reduced_beam_characteristics()
 
     # clean shutdown

diff --git a/tests/python/transformation.rst b/tests/python/transformation.rst
@@ -11,7 +11,7 @@ The beam has average energy 1 GeV.
 This tests that the t/s transforms are inverses of each other
 Specifically, in this test the :math:`t`- and :math:`s`-coordinates of the beam must differ substantially
 and the forward-inverse transformed coordinates must agree with the initial coordinates.
-That is, we require that ``to_fixed_s`` (``to_fixed_t`` (initial beam)) = initial beam.
+That is, we require that ``direction=CoordSystem.s`` (``direction=CoordSystem.t`` (initial beam)) = initial beam.
 
 
 Run