ImpactXParticleContainer.H

Go to the documentation of this file.

/* 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_PARTICLE_CONTAINER_H
#define IMPACTX_PARTICLE_CONTAINER_H
 
#include "ReferenceParticle.H"
#include "initialization/AmrCoreData_fwd.H"
 
#include <AMReX_AmrCoreFwd.H>
#include <AMReX_BaseFwd.H>
#include <AMReX_MultiFab.H>
#include <AMReX_ParIter.H>
#include <AMReX_Particles.H>
#include <AMReX_ParticleTile.H>
#include <AMReX_REAL.H>
 
#include <AMReX_IntVect.H>
#include <AMReX_Vector.H>
 
#include <list>
#include <optional>
#include <tuple>
#include <unordered_map>
 
 
namespace impactx
{
    enum class CoordSystem
    {
        s,  
        t   
    };

    struct RealSoA
    {
        enum
        {
            x,   
            y,   
            t,   
            px,  
            py,  
            pt,  
            sx,  
            sy,  
            sz,  
            qm,  
            w,   
            nattribs 
        };
 
        // at fixed t, the third component represents the position z, the 6th component represents the momentum in z
        enum {
            z = t,   
            pz = pt  
        };

        static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "spin_x", "spin_y", "spin_z", "qm", "weighting" };
        static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "spin_x", "spin_y", "spin_z", "qm", "weighting" };
        static_assert(names_s.size() == nattribs);
        static_assert(names_t.size() == nattribs);
    };

    struct IntSoA
    {
        enum
        {
            nattribs 
        };

        static constexpr std::initializer_list<char const *> names_s = {};
        static constexpr std::initializer_list<char const *> names_t = {};
        static_assert(names_s.size() == nattribs);
        static_assert(names_t.size() == nattribs);
    };

    class ParIterSoA
        : public amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs, amrex::PolymorphicArenaAllocator>
    {
    public:
        using amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs, amrex::PolymorphicArenaAllocator>::ParIterSoA;
 
        ParIterSoA (ContainerType& pc, int level);
 
        ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
 
        ContainerType * pc () { return m_pc; }
    };

    class ParConstIterSoA
        : public amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs, amrex::PolymorphicArenaAllocator>
    {
    public:
        using amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs, amrex::PolymorphicArenaAllocator>::ParConstIterSoA;
 
        ParConstIterSoA (ContainerType& pc, int level);
 
        ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
 
        ContainerType const * pc () const { return m_pc; }
    };

    class ImpactXParticleContainer
        : public amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs, amrex::PolymorphicArenaAllocator>
    {
    public:
        using iterator = impactx::ParIterSoA;

        using const_iterator = impactx::ParConstIterSoA;

        ImpactXParticleContainer (initialization::AmrCoreData* amr_core);

        virtual ~ImpactXParticleContainer() = default;

        void prepare ();

        void clear (bool keep_mass=false, bool keep_charge=false);

        void
        AddNParticles (
            amrex::Gpu::DeviceVector<amrex::ParticleReal> const & x,
            amrex::Gpu::DeviceVector<amrex::ParticleReal> const & y,
            amrex::Gpu::DeviceVector<amrex::ParticleReal> const & t,
            amrex::Gpu::DeviceVector<amrex::ParticleReal> const & px,
            amrex::Gpu::DeviceVector<amrex::ParticleReal> const & py,
            amrex::Gpu::DeviceVector<amrex::ParticleReal> const & pt,
            amrex::ParticleReal qm,
            std::optional<amrex::ParticleReal> bunch_charge = std::nullopt,
            std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> w = std::nullopt,
            std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sx = std::nullopt,
            std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sy = std::nullopt,
            std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sz = std::nullopt
        );

        void
        SetLostParticleContainer (ImpactXParticleContainer * lost_pc);
 
        ImpactXParticleContainer *
        GetLostParticleContainer ();

        void
        SetRefParticle (RefPart const & refpart);

        RefPart &
        GetRefParticle ();

        RefPart const &
        GetRefParticle () const;

        void SetRefParticleEdge ();

        int
        GetParticleShape () const { return m_particle_shape.value(); }

        void SetParticleShape ();

        void SetParticleShape (int order);

        std::tuple<
            amrex::ParticleReal, amrex::ParticleReal,
            amrex::ParticleReal, amrex::ParticleReal,
            amrex::ParticleReal, amrex::ParticleReal>
        MinAndMaxPositions ();

        std::tuple<
                amrex::ParticleReal, amrex::ParticleReal,
                amrex::ParticleReal, amrex::ParticleReal,
                amrex::ParticleReal, amrex::ParticleReal>
        MeanAndStdPositions ();

        void
        DepositCharge (std::unordered_map<int, amrex::MultiFab> & rho,
                       amrex::Vector<amrex::IntVect> const & ref_ratio);

        CoordSystem
        GetCoordSystem () const;

        void
        SetCoordSystem (CoordSystem coord_system);

        bool store_beam_moments = false;

        void
        record_beam_moments ();

        std::unordered_map<std::string, amrex::ParticleReal>
        beam_moments ();

        std::list<std::unordered_map<std::string, amrex::ParticleReal> >
        beam_moments_history () { return m_beam_moments; }

        void
        reset_beam_moments_history () { m_beam_moments.clear(); }
 
      private:

        RefPart m_refpart;

        std::optional<int> m_particle_shape;

        ImpactXParticleContainer* m_particles_lost = nullptr;

        CoordSystem m_coordsystem = CoordSystem::s;

        std::list<std::unordered_map<std::string, amrex::ParticleReal> > m_beam_moments;
 
    }; // ImpactXParticleContainer
 
} // namespace impactx
 
#endif // IMPACTX_PARTICLE_CONTAINER_H