diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp index 07d2cd7c4..4e8e9f683 100644 --- a/src/ImpactX.cpp +++ b/src/ImpactX.cpp @@ -183,7 +183,7 @@ namespace impactx // assuming that the distribution did not change // push all particles with external maps - Push(*m_particle_container, element_variant); + Push(*m_particle_container, element_variant, global_step); // just prints an empty newline at the end of the slice_step amrex::Print() << "\n"; diff --git a/src/particles/Push.H b/src/particles/Push.H index e8c0d1296..7665df7a6 100644 --- a/src/particles/Push.H +++ b/src/particles/Push.H @@ -20,11 +20,13 @@ namespace impactx { /** Push particles * - * @param pc container of the particles to push - * @param element_variant a single element to push the particles through + * @param[inout] pc container of the particles to push + * @param[inout] element_variant a single element to push the particles through + * @param[in] step global step for diagnostics */ void Push (ImpactXParticleContainer & pc, - KnownElements const & element_variant); + KnownElements & element_variant, + int step); } // namespace impactx diff --git a/src/particles/Push.cpp b/src/particles/Push.cpp index e610b6fb1..6523d95a9 100644 --- a/src/particles/Push.cpp +++ b/src/particles/Push.cpp @@ -10,56 +10,30 @@ #include "Push.H" #include -#include // for AMREX_RESTRICT -#include // for ParticleReal -#include +#include +#include namespace impactx { void Push (ImpactXParticleContainer & pc, - KnownElements const & element_variant) + KnownElements & element_variant, + int step) { - // performance profiling per element - std::string element_name; - std::visit([&element_name](auto&& element){ element_name = element.name; }, element_variant); - std::string const profile_name = "impactx::Push::" + element_name; - BL_PROFILE("impactx::Push"); - BL_PROFILE(profile_name); - - using namespace amrex::literals; // for _rt and _prt - - // preparing to access reference particle data: RefPart - RefPart & ref_part = pc.GetRefParticle(); - - // loop over refinement levels - int const nLevel = pc.finestLevel(); - for (int lev = 0; lev <= nLevel; ++lev) + // here we just access the element by its respective type + std::visit([&pc, step](auto&& element) { - // get simulation geometry information - //const amrex::Geometry& gm = this->Geom(lev); - //const auto prob_lo = gm.ProbLo(); - - // loop over all particle boxes - using ParIt = ImpactXParticleContainer::iterator; -#ifdef AMREX_USE_OMP -#pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) -#endif - for (ParIt pti(pc, lev); pti.isValid(); ++pti) { - // here we just access the element by its respective type - std::visit( - [&pti, &ref_part](auto element) { - // push reference particle in global coordinates - element(ref_part); - - // push beam particles relative to reference particle - element(pti, ref_part); - }, - element_variant - ); - } // end loop over all particle boxes - } // env mesh-refinement level loop + // performance profiling per element + std::string element_name; + element_name = element.name; + std::string const profile_name = "impactx::Push::" + element_name; + BL_PROFILE("impactx::Push"); + BL_PROFILE(profile_name); + + // push reference particle & all particles + element(pc, step); + }, element_variant); } } // namespace impactx diff --git a/src/particles/PushAll.H b/src/particles/PushAll.H new file mode 100644 index 000000000..2746058ce --- /dev/null +++ b/src/particles/PushAll.H @@ -0,0 +1,59 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_PUSH_ALL_H +#define IMPACTX_PUSH_ALL_H + +#include "particles/ImpactXParticleContainer.H" + + +namespace impactx +{ + /** Push all particles in a particle container. + * + * This element pushes first the reference particle, then all other particles. + * All particles are pushed independently with the same logic. + * Particles are relative to the reference particle. + * + * @param[in,out] pc particle container to push + * @param[in,out] element the beamline element + * @param[in] step global step for diagnostics + */ + template + void push_all ( + ImpactXParticleContainer & pc, + T_Element & element, + [[maybe_unused]] int step + ) + { + // preparing to access reference particle data: RefPart + RefPart & ref_part = pc.GetRefParticle(); + + // loop over refinement levels + int const nLevel = pc.finestLevel(); + for (int lev = 0; lev <= nLevel; ++lev) + { + // loop over all particle boxes + using ParIt = ImpactXParticleContainer::iterator; +#ifdef AMREX_USE_OMP +#pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) +#endif + for (ParIt pti(pc, lev); pti.isValid(); ++pti) { + // push reference particle in global coordinates + element(ref_part); + + // push beam particles relative to reference particle + element(pti, ref_part); + } // end loop over all particle boxes + } // env mesh-refinement level loop + } + +} // namespace impactx + +#endif // IMPACTX_PUSH_ALL_H diff --git a/src/particles/elements/None.H b/src/particles/elements/None.H index 9d10fdfd7..2ac2a8525 100644 --- a/src/particles/elements/None.H +++ b/src/particles/elements/None.H @@ -31,6 +31,14 @@ namespace impactx { } + /** Push all particles - nothing to do here */ + void operator() ( + ImpactXParticleContainer & /* pc */, + int /* step */ + ) { + // nothing to do + } + /** Push all particles - nothing to do here */ void operator() ( ImpactXParticleContainer::iterator & /* pti */, diff --git a/src/particles/elements/Programmable.H b/src/particles/elements/Programmable.H index f30148302..616334b26 100644 --- a/src/particles/elements/Programmable.H +++ b/src/particles/elements/Programmable.H @@ -26,9 +26,21 @@ namespace impactx static constexpr auto name = "Programmable"; using PType = ImpactXParticleContainer::ParticleType; - /** This element writes the particle beam out to openPMD data. + /** This element can be programmed */ - Programmable () {} + Programmable (amrex::ParticleReal ds=0.0, int nslice=1) + : m_ds(ds), m_nslice(nslice) + {} + + /** Push all particles relative to the reference particle + * + * @param[in,out] pc particle container to push + * @param[in] step global step for diagnostics + */ + void operator() ( + ImpactXParticleContainer & pc, + int step + ) const; /** Push all particles relative to the reference particle */ void operator() ( @@ -68,6 +80,7 @@ namespace impactx amrex::ParticleReal m_ds = 0.0; //! segment length in m int m_nslice = 1; //! number of slices used for the application of space charge + std::function m_push; //! hook for push of whole container std::function m_beam_particles; //! hook for beam particles std::function m_ref_particle; //! hook for reference particle }; diff --git a/src/particles/elements/Programmable.cpp b/src/particles/elements/Programmable.cpp index 84527aefa..2bb89bd72 100644 --- a/src/particles/elements/Programmable.cpp +++ b/src/particles/elements/Programmable.cpp @@ -9,10 +9,25 @@ */ #include "Programmable.H" +#include "particles/PushAll.H" namespace impactx { + void + Programmable::operator() ( + ImpactXParticleContainer & pc, + int step + ) const + { + if (m_push == nullptr) { + push_all(pc, *this, step); + } + else { + m_push(&pc, step); + } + } + void Programmable::operator() ( ImpactXParticleContainer::iterator & pti, @@ -20,6 +35,7 @@ namespace impactx ) const { if (m_beam_particles == nullptr) + // TODO: only if verbose mode is set amrex::AllPrint() << "Programmable element - all particles: NO HOOK\n"; else m_beam_particles(&pti, ref_part); @@ -29,6 +45,7 @@ namespace impactx Programmable::operator() (RefPart & ref_part) const { if (m_ref_particle == nullptr) + // TODO: only if verbose mode is set amrex::AllPrint() << "Programmable element - ref particles: NO HOOK\n"; else m_ref_particle(ref_part); diff --git a/src/particles/elements/mixin/beamoptic.H b/src/particles/elements/mixin/beamoptic.H index 1547494a5..4461558bd 100644 --- a/src/particles/elements/mixin/beamoptic.H +++ b/src/particles/elements/mixin/beamoptic.H @@ -11,10 +11,13 @@ #define IMPACTX_ELEMENTS_MIXIN_BEAMOPTIC_H #include "particles/ImpactXParticleContainer.H" +#include "particles/PushAll.H" -#include +#include // for AMREX_RESTRICT #include +#include + namespace impactx::elements { @@ -130,6 +133,21 @@ namespace detail template struct BeamOptic { + /** Push first the reference particle, then all other particles */ + void operator() ( + ImpactXParticleContainer & pc, + int step + ) + { + static_assert( + std::is_base_of_v, + "BeamOptic can only be used as a mixin class!" + ); + + T_Element& element = *static_cast(this); + push_all(pc, element, step); + } + /** This pushes the particles on a particle iterator tile or box. * * Particles are relative to the reference particle. @@ -140,7 +158,13 @@ namespace detail void operator() ( ImpactXParticleContainer::iterator & pti, RefPart & AMREX_RESTRICT ref_part - ) { + ) + { + static_assert( + std::is_base_of_v, + "BeamOptic can only be used as a mixin class!" + ); + T_Element& element = *static_cast(this); detail::push_all_particles(pti, ref_part, element); } diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 74acaf450..f944e44a1 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -88,7 +88,7 @@ void init_elements(py::module& m) .def_property_readonly("ds", &elements::Thin::ds) ; - // beam optics below + // beam optics py::class_(me, "ConstF") .def(py::init< @@ -148,7 +148,10 @@ void init_elements(py::module& m) ; py::class_(me, "Programmable") - .def(py::init<>(), + .def(py::init< + amrex::ParticleReal, + int>(), + py::arg("ds") = 0.0, py::arg("nslice") = 1, "A programmable beam optics element." ) .def_property("nslice", @@ -159,17 +162,26 @@ void init_elements(py::module& m) [](Programmable & p) { return p.ds(); }, [](Programmable & p, amrex::ParticleReal ds) { p.m_ds = ds; } ) + .def_property("push", + [](Programmable & p) { return p.m_push; }, + [](Programmable & p, + std::function new_hook + ) { p.m_push = new_hook; }, + "hook for push of whole container (pc, step)" + ) .def_property("beam_particles", [](Programmable & p) { return p.m_beam_particles; }, [](Programmable & p, std::function new_hook - ) { p.m_beam_particles = new_hook; } + ) { p.m_beam_particles = new_hook; }, + "hook for beam particles (pti, RefPart)" ) .def_property("ref_particle", [](Programmable & p) { return p.m_ref_particle; }, [](Programmable & p, std::function new_hook - ) { p.m_ref_particle = new_hook; } + ) { p.m_ref_particle = new_hook; }, + "hook for reference particle (RefPart)" ) ;