From b8169b40520191b2ad1ae34aeb9e52f600943fc9 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sat, 6 Jan 2024 21:54:34 +0100 Subject: [PATCH] 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. --- src/ImpactX.cpp | 4 ++-- src/particles/ImpactXParticleContainer.H | 20 +++++++++++++++++++ src/particles/ImpactXParticleContainer.cpp | 12 +++++++++++ .../transformation/CoordinateTransformation.H | 12 ++--------- .../CoordinateTransformation.cpp | 15 +++++++++++--- src/python/ImpactXParticleContainer.cpp | 10 ++++++++++ src/python/transformation.cpp | 7 ++----- tests/python/test_transformation.py | 20 +++++++++++++++---- tests/python/transformation.rst | 2 +- 9 files changed, 77 insertions(+), 25 deletions(-) diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp index 7e948c4f0..895d318a1 100644 --- 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 index 4d4f7dee5..4c1fd480c 100644 --- 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 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 index e18d53341..b82925eac 100644 --- 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 get_RealAoS_names () { diff --git a/src/particles/transformation/CoordinateTransformation.H b/src/particles/transformation/CoordinateTransformation.H index c104bb9a3..7a5459b64 100644 --- 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 index dbf143698..b52e86afe 100644 --- 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 index e9f25f9e3..3b6fc3f57 100644 --- 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_(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 index 1eaef9cb5..222d66602 100644 --- 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_(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 index 60bba3b5f..d44ec073e 100644 --- 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 index 17aaae8a8..7d9b988c0 100644 --- 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