particles.H

Go to the documentation of this file.

/* Copyright 2022-2025 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, Chad Mitchell
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_TRACKING_PARTICLES_H
#define IMPACTX_TRACKING_PARTICLES_H
 
#include "elements/All.H"
#include "particles/ImpactXParticleContainer.H"
#include "tracking/TrackingDirection.H"
#include "tracking/TrackingState.H"
 
#include <AMReX_BLProfiler.H>
#include <AMReX_ParmParse.H>
#include <AMReX_Print.H>
#include <AMReX_REAL.H>
 
#include <list>
#include <string>
#include <variant>
 
 
namespace impactx
{
    template <typename HookFn, typename CollectiveKicks, typename ElementPush>
    void track_lattice_particles (
        std::list<elements::KnownElements> & lattice,
        ImpactXParticleContainer & pc,
        TrackingState & tracking_state,
        HookFn const & call_hook,
        CollectiveKicks && collective_kicks,
        ElementPush && element_push
    )
    {
        bool const forward = tracking_state.forward();
 
        // periods through the lattice (turns/cycles) and verbosity
        int num_periods = 1;
        amrex::ParmParse("lattice").queryAddWithParser("periods", num_periods);
        int verbose = 1;
        amrex::ParmParse("impactx").queryAddWithParser("verbose", verbose);
 
        int & step = tracking_state.m_step;
 
        // visit one element for all of its space-charge slices
        auto visit_element = [&] (
            elements::KnownElements & element_variant,
            int period
        )
        {
            if (forward)
            {
                // update element edge of the reference particle
                pc.SetRefParticleEdge();
 
                // book-keeping: currently visited element reference
                tracking_state.m_element = &element_variant;
 
                // optional, user-defined function call
                call_hook("before_element");
            }
            else
            {
                // back-tracking: invert this element's map in place so the element
                // transport below applies the inverse map.
                std::visit([](auto && element){ element.reverse(); }, element_variant);
            }
 
            // number of slices used for the application of space charge
            int nslice = 1;
            amrex::ParticleReal slice_ds = 0; // in meters
            std::visit([&nslice, &slice_ds](auto && element) {
                nslice = element.nslice();
                slice_ds = element.ds() / nslice;
            }, element_variant);
 
            // sub-steps for space charge within the element
            for (int s = 0; s < nslice; ++s)
            {
                BL_PROFILE("ImpactX::evolve::slice_step");
                int const slice_step = forward ? s : (nslice - 1 - s);
 
                if (forward)
                {
                    step++;
                    if (verbose > 0) {
                        amrex::Print() << "\n++++ Starting step=" << step
                                       << " slice_step=" << slice_step;
                    }
 
                    // optional, user-defined function call
                    call_hook("before_slice");
                }
 
                // TODO: Strang-split here (#846 / #920)
                if (forward)
                {
                    collective_kicks(element_variant, slice_ds);
                    element_push(element_variant, step, period);
                }
                else
                {
                    element_push(element_variant, step, period);
                    collective_kicks(element_variant, slice_ds);
                }
            }
 
            if (forward)
            {
                // optional, user-defined function call
                call_hook("after_element");
            }
            else
            {
                // restore the forward element map (reverse() is an involution)
                std::visit([](auto && element){ element.reverse(); }, element_variant);
            }
        };
 
        // loop over channel periods or ring turns
        for (int p = 0; p < num_periods; ++p)
        {
            int const period = forward ? p : (num_periods - 1 - p);
 
            if (forward)
            {
                // book-keeping update
                tracking_state.m_period = period;
                tracking_state.m_element = &lattice.front();
 
                // optional, user-defined function call
                call_hook("before_period");
            }
 
            // loop over lattice elements
            if (forward)
            {
                for (auto & element_variant : lattice)
                {
                    visit_element(element_variant, period);
                }
            }
            else
            {
                for (auto it = lattice.rbegin(); it != lattice.rend(); ++it)
                {
                    visit_element(*it, period);
                }
            }
 
            // optional, user-defined function call
            if (forward) { call_hook("after_period"); }
        }
 
        if (forward)
        {
            // avoid dangling references if users manipulate the lattice
            tracking_state.set_no_element();
        }
    }
} // namespace impactx
 
#endif // IMPACTX_TRACKING_PARTICLES_H