WarpXParticleContainer.H

Go to the documentation of this file.

/* Copyright 2019-2020 Andrew Myers, Axel Huebl, David Grote
 * Jean-Luc Vay, Junmin Gu, Luca Fedeli
 * Maxence Thevenet, Remi Lehe, Revathi Jambunathan
 * Weiqun Zhang, Yinjian Zhao
 *
 * This file is part of WarpX.
 *
 * License: BSD-3-Clause-LBNL
 */
#ifndef WARPX_WarpXParticleContainer_H_
#define WARPX_WarpXParticleContainer_H_
 
#include "WarpXParticleContainer_fwd.H"
#include "FieldSolver/ImplicitSolvers/ImplicitOptions.H"
 
#include "Initialization/PlasmaInjector.H"
#include "Particles/Algorithms/KineticEnergy.H"
#include "Particles/ParticleBoundaries.H"
#include "SpeciesPhysicalProperties.H"
 
#ifdef WARPX_QED
#    include "ElementaryProcess/QEDInternals/BreitWheelerEngineWrapper_fwd.H"
#    include "ElementaryProcess/QEDInternals/QuantumSyncEngineWrapper_fwd.H"
#endif
#include "MultiParticleContainer_fwd.H"
 
#include "Particles/Deposition/TemperatureDepositionTypes.H"
#include "Particles/Deposition/VarianceAccumulationBuffer.H"
 
#include <ablastr/fields/MultiFabRegister.H>
 
#include <AMReX_Array.H>
#include <AMReX_FArrayBox.H>
#include <AMReX_Geometry.H>
#include <AMReX_GpuAllocators.H>
#include <AMReX_GpuContainers.H>
#include <AMReX_INT.H>
#include <AMReX_ParIter.H>
#include <AMReX_Particles.H>
#include <AMReX_Random.H>
#include <AMReX_REAL.H>
#include <AMReX_StructOfArrays.H>
#include <AMReX_Vector.H>
 
#include <AMReX_BaseFwd.H>
#include <AMReX_AmrCoreFwd.H>
 
#include <array>
#include <iosfwd>
#include <map>
#include <memory>
#include <string>
#include <utility>

struct PIdx
{
#if defined(WARPX_DIM_1D_Z)
    enum                          { z ,  w ,  ux ,  uy ,  uz,                      nattribs};
    static constexpr auto names = {"z", "w", "ux", "uy", "uz"};
#elif defined(WARPX_DIM_XZ)
    enum                          { x ,  z ,  w ,  ux ,  uy ,  uz,                 nattribs};
    static constexpr auto names = {"x", "z", "w", "ux", "uy", "uz"};
#elif defined(WARPX_DIM_RZ)
    enum                          { r ,  z ,  w ,  ux ,  uy ,  uz ,  theta,        nattribs};
    static constexpr auto names = {"r", "z", "w", "ux", "uy", "uz", "theta"};
#elif defined(WARPX_DIM_RCYLINDER)
    enum                          { r ,  w ,  ux ,  uy ,  uz ,  theta,             nattribs};
    static constexpr auto names = {"r", "w", "ux", "uy", "uz", "theta"};
#elif defined(WARPX_DIM_RSPHERE)
    enum                          { r ,  w ,  ux ,  uy ,  uz ,  theta ,  phi,      nattribs};
    static constexpr auto names = {"r", "w", "ux", "uy", "uz", "theta", "phi"};
#elif defined(WARPX_DIM_3D)
    enum                          { x ,  y,   z,   w,   ux,   uy,   uz,            nattribs};
    static constexpr auto names = {"x", "y", "z", "w", "ux", "uy", "uz"};
#else
  #error WarpX dimension is not correctly set!
#endif
 
    static_assert(names.size() == nattribs);
};
 
 
struct IntIdx {
    enum
        {
            nattribs 
        };
 
    static constexpr std::initializer_list<char const *> names = {};
};
 
struct DiagIdx
{
    enum {
        w = 0,
        x, y, z, ux, uy, uz,
        nattribs
    };
};
 
class WarpXParIter
    : public amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PolymorphicArenaAllocator>
{
public:
    using amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PolymorphicArenaAllocator>::ParIterSoA;
 
    WarpXParIter (ContainerType& pc, int level);
 
    WarpXParIter (ContainerType& pc, int level, amrex::MFItInfo& info);
 
    [[nodiscard]] const std::array<RealVector, PIdx::nattribs>& GetAttribs () const
    {
        return GetStructOfArrays().GetRealData();
    }
 
    [[nodiscard]] std::array<RealVector, PIdx::nattribs>& GetAttribs ()
    {
        return GetStructOfArrays().GetRealData();
    }
 
    [[nodiscard]] const RealVector& GetAttribs (int comp) const
    {
        return GetStructOfArrays().GetRealData(comp);
    }
 
    [[nodiscard]] RealVector& GetAttribs (int comp)
    {
        return GetStructOfArrays().GetRealData(comp);
    }
 
    [[nodiscard]] const IntVector& GetiAttribs (int comp) const
    {
        return GetStructOfArrays().GetIntData(comp);
    }
 
    [[nodiscard]] IntVector& GetiAttribs (int comp)
    {
        return GetStructOfArrays().GetIntData(comp);
    }
 
    [[nodiscard]] const RealVector& GetAttribs (const std::string& name) const
    {
        return GetStructOfArrays().GetRealData(name);
    }
 
    [[nodiscard]] RealVector& GetAttribs (const std::string& name)
    {
        return GetStructOfArrays().GetRealData(name);
    }
 
    [[nodiscard]] const IntVector& GetiAttribs (const std::string& name) const
    {
        return GetStructOfArrays().GetIntData(name);
    }
 
    [[nodiscard]] IntVector& GetiAttribs (const std::string& name)
    {
        return GetStructOfArrays().GetIntData(name);
    }
};

class WarpXParticleContainer
    : public amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, amrex::PolymorphicArenaAllocator>
{
public:
    friend MultiParticleContainer;
 
    using Base = amrex::ParticleContainerPureSoA<PIdx::nattribs, 0, amrex::PolymorphicArenaAllocator>;
 
    // amrex::StructOfArrays with DiagIdx::nattribs amrex::ParticleReal components
    // and 0 int components for the particle data.
    using DiagnosticParticleData = amrex::StructOfArrays<DiagIdx::nattribs, 0>;
    // DiagnosticParticles is a vector, with one element per MR level.
    // DiagnosticParticles[lev] is typically a key-value pair where the key is
    // a pair [grid_index, tile_index], and the value is the corresponding
    // DiagnosticParticleData (see above) on this tile.
    using DiagnosticParticles = amrex::Vector<std::map<std::pair<int, int>, DiagnosticParticleData> >;
 
    WarpXParticleContainer (amrex::AmrCore* amr_core, int ispecies, const std::string& name);
    ~WarpXParticleContainer() override = default;
 
    // Move and copy operations
    WarpXParticleContainer(const WarpXParticleContainer&)            = delete;
    WarpXParticleContainer& operator=(const WarpXParticleContainer&) = delete;
    WarpXParticleContainer(WarpXParticleContainer&&)                 = default;
    WarpXParticleContainer& operator=(WarpXParticleContainer&&)      = default;
 
    virtual void InitData () = 0;
 
    virtual void InitIonizationModule () {}
 
    /*
     * \brief Virtual function that returns a pointer to the plasma injector,
     *        for derived classes that define one (PhysicalParticleContainer).
     */
    virtual PlasmaInjector* GetPlasmaInjector (const int /*i*/) { return nullptr; }

    virtual void Evolve (ablastr::fields::MultiFabRegister& fields,
                         int lev,
                         const std::string& current_fp_string,
                         amrex::Real t, amrex::Real dt, SubcyclingHalf subcycling_half=SubcyclingHalf::None, bool skip_deposition=false,
                         PositionPushType position_push_type=PositionPushType::Full,
                         MomentumPushType momentum_push_type=MomentumPushType::Full,
                         ImplicitOptions const * implicit_options = nullptr) = 0;

    virtual void FinishImplicitParticleUpdate (
        ablastr::fields::MultiFabRegister& fields,
        int lev, amrex::Real t, amrex::Real dt);

    virtual void DepositMassMatrices (ablastr::fields::MultiFabRegister& fields,
                                      int lev,
                                      amrex::Real dt)
    {
        amrex::ignore_unused(fields,lev,dt);
    }
 
    virtual void PostRestart () = 0;
 
    virtual void AllocData ();

    virtual void DefaultInitializeRuntimeAttributes (
        typename ContainerLike<amrex::PolymorphicArenaAllocator>::ParticleTileType& pinned_tile,
        int n_external_attr_real,
        int n_external_attr_int) = 0;

    void PushX (         amrex::Real dt);
    void PushX (int lev, amrex::Real dt);

    virtual void PushP (int lev, amrex::Real dt,
                        const amrex::MultiFab& Ex,
                        const amrex::MultiFab& Ey,
                        const amrex::MultiFab& Ez,
                        const amrex::MultiFab& Bx,
                        const amrex::MultiFab& By,
                        const amrex::MultiFab& Bz,
                        MomentumPushType momentum_push_type = MomentumPushType::Full) = 0;

    void DepositCurrent (ablastr::fields::MultiLevelVectorField const & J,
                         amrex::Real dt, amrex::Real relative_time,
                         PushType push_type = PushType::Explicit);

    void DepositCurrent (const std::string& mf_name, int lev,
                         amrex::Real dt, amrex::Real relative_time);

    void DepositCharge (const ablastr::fields::MultiLevelScalarField& rho,
                        bool local = false, bool reset = false,
                        bool apply_boundary_and_scale_volume = false,
                        bool interpolate_across_levels = true,
                        int icomp = 0);
    void DepositCharge (amrex::MultiFab* rho, int lev,
                        bool local = false, bool reset = false,
                        bool apply_boundary_and_scale_volume = false,
                        int icomp = 0);
 
    std::unique_ptr<amrex::MultiFab> GetChargeDensity(int lev, bool local = false);
 
    virtual void DepositCharge (WarpXParIter& pti,
                               RealVector const & wp,
                               const int* ion_lev,
                               amrex::MultiFab* rho,
                               int icomp,
                               long offset,
                               long np_to_deposit,
                               int thread_num,
                               int lev,
                               int depos_lev);
 
    virtual void DepositCurrent (WarpXParIter& pti,
                                RealVector const & wp,
                                RealVector const & uxp,
                                RealVector const & uyp,
                                RealVector const & uzp,
                                int const* ion_lev,
                                amrex::MultiFab* jx,
                                amrex::MultiFab* jy,
                                amrex::MultiFab* jz,
                                long offset,
                                long np_to_deposit,
                                int thread_num,
                                int lev,
                                int depos_lev,
                                amrex::Real dt,
                                amrex::Real relative_time,
                                PushType push_type);
    virtual void AccumulateVelocitiesAndComputeTemperature (
        [[maybe_unused]]ablastr::fields::MultiLevelVectorField const & T_vf,
        [[maybe_unused]]amrex::Real relative_time)
    {}
 
    virtual void DepositTemperature (
        [[maybe_unused]]WarpXParIter& pti,
        [[maybe_unused]]RealVector const & wp,
        [[maybe_unused]]RealVector const & uxp,
        [[maybe_unused]]RealVector const & uyp,
        [[maybe_unused]]RealVector const & uzp,
        [[maybe_unused]]amrex::MultiFab * const Tx,
        [[maybe_unused]]amrex::MultiFab * const Ty,
        [[maybe_unused]]amrex::MultiFab * const Tz,
        [[maybe_unused]]long const offset,
        [[maybe_unused]]long const np_to_deposit,
        [[maybe_unused]]int const thread_num,
        [[maybe_unused]]const int lev,
        [[maybe_unused]]int const depos_lev,
        [[maybe_unused]]amrex::Real const relative_time,
        [[maybe_unused]]PushType push_type,
        [[maybe_unused]]const warpx::particles::deposition::TemperatureDepositionType type,
        [[maybe_unused]]const warpx::particles::deposition::TemperatureDepositionPass pass)
    {}
 
    void DepositTotalNGPTemperature (int lev);
    void CalculateNuei(amrex::MultiFab & species_nuei,
                       WarpXParticleContainer const & electron_species, int lev);
    std::unique_ptr<amrex::MultiFab> GetDebyeLength (int lev);
 
    void DepositNumberDensity (amrex::MultiFab* number_density, int lev);
    std::unique_ptr<amrex::MultiFab> GetNumberDensity (int lev);
 
    std::unique_ptr<amrex::MultiFab> GetPlasmaFrequency (int lev);
 
    void DepositMassMatrices (WarpXParIter& pti,
                        const RealVector& wp,
                        const RealVector& uxp,
                        const RealVector& uyp,
                        const RealVector& uzp,
                              amrex::MultiFab* Sxx,
                              amrex::MultiFab* Sxy,
                              amrex::MultiFab* Sxz,
                              amrex::MultiFab* Syx,
                              amrex::MultiFab* Syy,
                              amrex::MultiFab* Syz,
                              amrex::MultiFab* Szx,
                              amrex::MultiFab* Szy,
                              amrex::MultiFab* Szz,
                        const amrex::FArrayBox* bxfab,
                        const amrex::FArrayBox* byfab,
                        const amrex::FArrayBox* bzfab,
                              long offset,
                              long np_to_deposit,
                              int thread_num,
                              int lev,
                              int depos_lev,
                              amrex::Real dt);
 
    // If particles start outside of the domain, ContinuousInjection
    // makes sure that they are initialized when they enter the domain, and
    // NOT before. Virtual function, overriden by derived classes.
    // Current status:
    // PhysicalParticleContainer: implemented.
    // LaserParticleContainer: implemented.
    // RigidInjectedParticleContainer: not implemented.
    virtual void ContinuousInjection(const amrex::RealBox& /*injection_box*/) {}

    virtual void UpdateAntennaPosition(const amrex::Real /*dt*/) {}
 
    bool doContinuousInjection() const {return do_continuous_injection;}
 
    // Inject a continuous flux of particles from a defined plane
    virtual void ContinuousFluxInjection(amrex::Real /*t*/, amrex::Real /*dt*/) {}
 
    int getSpeciesId() const {return species_id;}

    std::pair<amrex::ParticleReal, amrex::ParticleReal> sumParticleWeightAndEnergy (bool local = false) const;
    amrex::ParticleReal sumParticleCharge (bool local = false) const;
    amrex::ParticleReal sumParticleWeight (bool local = false) const;
    amrex::ParticleReal sumParticleEnergy (bool local = false) const;
 
    std::array<amrex::ParticleReal, 3> meanParticleVelocity(bool local = false);
 
    amrex::ParticleReal maxParticleDtInv(bool local = false);

    void TransformMomentumToCurvilinear (bool forward);

    void AddNParticles (int lev, long n,
                        amrex::Vector<amrex::ParticleReal> const & x,
                        amrex::Vector<amrex::ParticleReal> const & y,
                        amrex::Vector<amrex::ParticleReal> const & z,
                        amrex::Vector<amrex::ParticleReal> const & ux,
                        amrex::Vector<amrex::ParticleReal> const & uy,
                        amrex::Vector<amrex::ParticleReal> const & uz,
                        int nattr_real,
                        amrex::Vector<amrex::Vector<amrex::ParticleReal>> const & attr_real,
                        int nattr_int, amrex::Vector<amrex::Vector<int>> const & attr_int,
                        int uniqueparticles, amrex::Long id=-1);

    void deleteInvalidParticles ();
 
    virtual void ReadHeader (std::istream& is) = 0;
 
    virtual void WriteHeader (std::ostream& os) const = 0;
 
    static void ReadParameters ();
 
    static void BackwardCompatibility ();

    void ApplyBoundaryConditions ();
 
    bool do_splitting = false;
 
    int do_not_deposit = 0;
    bool initialize_self_fields = false;
    amrex::Real self_fields_required_precision = amrex::Real(1.e-11);
    amrex::Real self_fields_absolute_tolerance = amrex::Real(0.0);
    int self_fields_max_iters = 200;
    int self_fields_verbosity = 2;
 
    // External fields added to particle fields.
    amrex::Vector<amrex::ParticleReal> m_B_external_particle;
    amrex::Vector<amrex::ParticleReal> m_E_external_particle;

    amrex::Real m_current_injection_position = 0.;
 
    // split along diagonals (0) or axes (1)
    int split_type = 0;

    void SetDoBackTransformedParticles (const bool do_back_transformed_particles) {
        m_do_back_transformed_particles = do_back_transformed_particles;
    }
 
    bool GetDoBackTransformedParticles () const {
        return m_do_back_transformed_particles;
    }
 
    amrex::ParticleReal getCharge () const {return m_charge;}
 
    amrex::ParticleReal getMass () const {return m_mass;}

    amrex::ParticleReal getBoundaryThermalVelocity () const
    {
        return static_cast<amrex::ParticleReal>(m_boundary_conditions.data.m_uth);
    }
 
    int DoFieldIonization() const { return do_field_ionization; }
 
#ifdef WARPX_QED
    //Species for which QED effects are relevant should override these methods
    virtual bool has_quantum_sync() const {return false;}
    virtual bool has_breit_wheeler() const {return false;}
    virtual bool has_virtual_photons() const {return false;}
    virtual bool has_virtual_photons_beam_size_effect() const {return false;}
    virtual int getVirtualPhotonSpeciesIndex() const { return -1; }
    int DoQED() const { return has_quantum_sync() || has_breit_wheeler(); }
#else
    int DoQED() const { return false; }
#endif
 
    /* \brief This function tests if the current species
    *  is of a given PhysicalSpecies (specified as a template parameter).
    * @tparam PhysSpec the PhysicalSpecies to test against
    * @return the result of the test
    */
    template<PhysicalSpecies PhysSpec>
    bool AmIA () const noexcept {return (physical_species == PhysSpec);}

    std::string getSpeciesTypeName () const {return species::get_name(physical_species);}

    virtual void resample (const amrex::Vector<amrex::Geometry>& /*geom*/, const int /*timestep*/, bool /*verbose*/) {}

    void defineAllParticleTiles () noexcept;
 
    virtual std::vector<std::string> getUserIntAttribs () const { return std::vector<std::string>{}; }
 
    virtual std::vector<std::string> getUserRealAttribs () const { return std::vector<std::string>{}; }
 
    virtual amrex::Vector< amrex::Parser* > getUserIntAttribParser () const { return amrex::Vector<amrex::Parser*>{}; }
 
    virtual amrex::Vector< amrex::Parser* > getUserRealAttribParser () const { return amrex::Vector<amrex::Parser*>{}; }
 
#ifdef WARPX_QED
    virtual BreitWheelerEngine* get_breit_wheeler_engine_ptr () const {return nullptr;}
    virtual QuantumSynchrotronEngine* get_quantum_sync_engine_ptr () const {return nullptr;}
#endif
 
    void setDoNotPush (bool flag) { do_not_push = flag; }
    void setDoNotGather (int flag) { do_not_gather = flag; }
    void setDoNotDeposit (int flag) { do_not_deposit = flag; }
 
    bool HasiAttrib (const std::string& name)
    {
        std::vector<std::string> const & int_names = GetIntSoANames();
        auto const pos = std::find(int_names.begin(), int_names.end(), name);
        return (pos != int_names.end());
    }
 
    std::string const & getName() { return species_name; }
 
protected:
    int species_id;
    std::string species_name;
 
    amrex::ParticleReal m_charge;
    amrex::ParticleReal m_mass;
    PhysicalSpecies physical_species = PhysicalSpecies::unspecified;
 
    // Controls boundaries for particles exiting the domain
    ParticleBoundaries m_boundary_conditions;
 
    // Controls whether particle trajectoris are cropped when particles
    // cross a PEC boundary
    bool m_crop_on_PEC_boundary = false;

    bool m_deposit_on_main_grid = false;

    bool m_gather_from_main_grid = false;
 
    bool do_not_push = false;
    int do_not_gather = 0;
 
    // Whether to allow particles outside of the simulation domain to be
    // initialized when they enter the domain.
    // This is currently required because continuous injection does not
    // support all features allowed by direct injection.
    int do_continuous_injection = 0;
 
    int do_field_ionization = 0;
    int do_adk_correction = 0;
    int ionization_product;
    std::string ionization_product_name;
    int ion_atomic_number;
    int ionization_initial_level = 0;
    amrex::Gpu::DeviceVector<amrex::Real> ionization_energies;
    amrex::Gpu::DeviceVector<amrex::Real> adk_power;
    amrex::Gpu::DeviceVector<amrex::Real> adk_prefactor;
    amrex::Gpu::DeviceVector<amrex::Real> adk_exp_prefactor;
    amrex::Gpu::DeviceVector<amrex::Real> adk_correction_factors;
    std::string physical_element;
 
    int do_resampling = 0;

    bool m_do_back_transformed_particles = false;
 
#ifdef WARPX_QED
    //Species can receive a shared pointer to a QED engine (species for
    //which this is relevant should override these functions)
    virtual void
    set_breit_wheeler_engine_ptr(const std::shared_ptr<BreitWheelerEngine>&){}
    virtual void
    set_quantum_sync_engine_ptr(const std::shared_ptr<QuantumSynchrotronEngine>&){}
 
    int m_qed_breit_wheeler_ele_product;
    std::string m_qed_breit_wheeler_ele_product_name;
    int m_qed_breit_wheeler_pos_product;
    std::string m_qed_breit_wheeler_pos_product_name;
    int m_qed_quantum_sync_phot_product;
    std::string m_qed_quantum_sync_phot_product_name;
    int m_qed_virtual_photon_species;
    std::string m_qed_virtual_photon_species_name;
    amrex::Real m_qed_virtual_photon_min_energy;
#endif
    amrex::Vector<amrex::FArrayBox> local_rho;
    amrex::Vector<amrex::FArrayBox> local_jx;
    amrex::Vector<amrex::FArrayBox> local_jy;
    amrex::Vector<amrex::FArrayBox> local_jz;
 
    // local containers for mass matrix deposit used by implicit solvers
    amrex::Vector<amrex::FArrayBox> local_Sxx;
    amrex::Vector<amrex::FArrayBox> local_Sxy;
    amrex::Vector<amrex::FArrayBox> local_Sxz;
    amrex::Vector<amrex::FArrayBox> local_Syx;
    amrex::Vector<amrex::FArrayBox> local_Syy;
    amrex::Vector<amrex::FArrayBox> local_Syz;
    amrex::Vector<amrex::FArrayBox> local_Szx;
    amrex::Vector<amrex::FArrayBox> local_Szy;
    amrex::Vector<amrex::FArrayBox> local_Szz;
 
public:
 
    int getIonizationInitialLevel () const noexcept {return ionization_initial_level;}
 
    // By default this returns false, can be overridden in PysicalParticleContainer
    virtual bool getTemperatureDepositionFlag () const noexcept {return false;}
 
private:
    void particlePostLocate(ParticleType& p, const amrex::ParticleLocData& pld, int lev) override;
 
};
 
#endif