DefaultInitialization.H

Go to the documentation of this file.

/* Copyright 2019-2020 Andrew Myers, Axel Huebl,
 * Maxence Thevenet
 *
 * This file is part of WarpX.
 *
 * License: BSD-3-Clause-LBNL
 */
#ifndef WARPX_DEFAULTINITIALIZATION_H_
#define WARPX_DEFAULTINITIALIZATION_H_
 
#include <WarpX.H>
#ifdef WARPX_QED
#   include "Particles/ElementaryProcess/QEDInternals/BreitWheelerEngineWrapper.H"
#   include "Particles/ElementaryProcess/QEDInternals/QuantumSyncEngineWrapper.H"
#endif
 
#include <AMReX_GpuContainers.H>
#include <AMReX_REAL.H>
 
#include <cmath>
#include <map>
#include <string>

enum struct InitializationPolicy {Zero=0, One, RandomExp};

static std::map<std::string, InitializationPolicy> initialization_policies = {
    {"w",     InitializationPolicy::Zero },
    {"ux",    InitializationPolicy::Zero },
    {"uy",    InitializationPolicy::Zero },
    {"uz",    InitializationPolicy::Zero },
#ifdef WARPX_DIM_RZ
    {"theta", InitializationPolicy::Zero},
#endif
 
#ifdef WARPX_QED
    {"opticalDepthBW",   InitializationPolicy::RandomExp},
    {"opticalDepthQSR",   InitializationPolicy::RandomExp}
#endif
 
};
 
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
amrex::ParticleReal initializeRealValue (const InitializationPolicy policy, amrex::RandomEngine const& engine) noexcept
{
    switch (policy) {
        case InitializationPolicy::Zero : return 0.0;
        case InitializationPolicy::One  : return 1.0;
        case InitializationPolicy::RandomExp : {
            return -std::log(amrex::Random(engine));
        }
        default : {
            amrex::Abort("Initialization Policy not recognized");
            return 1.0;
        }
    }
}
 
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
int initializeIntValue (const InitializationPolicy policy) noexcept
{
    switch (policy) {
        case InitializationPolicy::Zero : return 0;
        case InitializationPolicy::One  : return 1;
        default : {
            amrex::Abort("Initialization Policy not recognized");
            return 1;
        }
    }
}
 
namespace ParticleCreation {

template <typename PTile, typename DstPC>
void DefaultInitializeRuntimeAttributes (PTile& ptile,
                                         const DstPC& pc,
                                         int start, int stop,
                                         const int n_external_attr_real = 0,
                                         const int n_external_attr_int = 0
#ifdef WARPX_QED
                                         , const bool do_qed_comps = false )
#else
                                         )
#endif
{
        using namespace amrex::literals;
 
        const std::vector<std::string>& user_real_attribs = pc.getUserRealAttribs();
        const std::vector<std::string>& user_int_attribs  = pc.getUserIntAttribs();
        const std::vector<std::string>& particle_comps  = pc.GetRealSoANames();
        const std::vector<std::string>& particle_icomps = pc.GetIntSoANames();
        const std::vector<amrex::Parser*>& user_real_attrib_parser = pc.getUserRealAttribParser();
        const std::vector<amrex::Parser*>& user_int_attrib_parser  = pc.getUserIntAttribParser();
#ifdef WARPX_QED
        BreitWheelerEngine*       p_bw_engine = pc.get_breit_wheeler_engine_ptr();
        QuantumSynchrotronEngine* p_qs_engine = pc.get_quantum_sync_engine_ptr();
#endif
        const int ionization_initial_level = pc.getIonizationInitialLevel();
 
        // Preparing data needed for user defined attributes
        const auto n_user_real_attribs = static_cast<int>(user_real_attribs.size());
        const auto n_user_int_attribs = static_cast<int>(user_int_attribs.size());
        const auto get_position = GetParticlePosition<PIdx>(ptile);
        const auto soa = ptile.getParticleTileData();
        const amrex::ParticleReal* AMREX_RESTRICT ux = soa.m_rdata[PIdx::ux];
        const amrex::ParticleReal* AMREX_RESTRICT uy = soa.m_rdata[PIdx::uy];
        const amrex::ParticleReal* AMREX_RESTRICT uz = soa.m_rdata[PIdx::uz];
        constexpr int lev = 0;
        const amrex::Real t = WarpX::GetInstance().gett_new(lev);
 
        // Check: does the particle data reside on CPU or GPU?
        amrex::Arena const * const pc_arena = pc.arena();
        const bool run_on_gpu =
            pc_arena->isManaged() || pc_arena->isDevice();
 
        // Initialize the last NumRuntimeRealComps() - n_external_attr_real runtime real attributes
        for (int j = PIdx::nattribs + n_external_attr_real; j < ptile.NumRealComps() ; ++j)
        {
            auto attr_ptr = ptile.GetStructOfArrays().GetRealData(j).data();
#ifdef WARPX_QED
            // Current runtime comp is quantum synchrotron optical depth
            auto const it_qsr = std::find(particle_comps.begin(), particle_comps.end(), "opticalDepthQSR");
            if (it_qsr != particle_comps.end() &&
                std::distance(particle_comps.begin(), it_qsr) == j)
            {
                if (!do_qed_comps) { continue; }
                const QuantumSynchrotronGetOpticalDepth quantum_sync_get_opt =
                                                p_qs_engine->build_optical_depth_functor();
                // If the particle tile was allocated in a memory pool that can run on GPU, launch GPU kernel
                if (run_on_gpu) {
                        amrex::ParallelForRNG(stop - start,
                                              [=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept {
                                                  const int ip = i + start;
                                                  attr_ptr[ip] = quantum_sync_get_opt(engine);
                                              });
                // Otherwise (e.g. particle tile allocated in pinned memory), run on CPU
                } else {
                    for (int ip = start; ip < stop; ++ip) {
#ifdef AMREX_USE_GPU
                        attr_ptr[ip] = quantum_sync_get_opt(amrex::RandomEngine{nullptr});
#else
                        attr_ptr[ip] = quantum_sync_get_opt(amrex::RandomEngine{});
#endif
                    }
                }
            }
 
            // Current runtime comp is Breit-Wheeler optical depth
            auto const it_bw = std::find(particle_comps.begin(), particle_comps.end(), "opticalDepthBW");
            if (it_bw != particle_comps.end() &&
                std::distance(particle_comps.begin(), it_bw) == j)
            {
                if (!do_qed_comps) { continue; }
                const BreitWheelerGetOpticalDepth breit_wheeler_get_opt =
                                                p_bw_engine->build_optical_depth_functor();;
                // If the particle tile was allocated in a memory pool that can run on GPU, launch GPU kernel
                if (run_on_gpu) {
                        amrex::ParallelForRNG(stop - start,
                                              [=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept {
                                                  const int ip = i + start;
                                                  attr_ptr[ip] = breit_wheeler_get_opt(engine);
                                              });
                // Otherwise (e.g. particle tile allocated in pinned memory), run on CPU
                } else {
                    for (int ip = start; ip < stop; ++ip) {
#ifdef AMREX_USE_GPU
                        attr_ptr[ip] = breit_wheeler_get_opt(amrex::RandomEngine{nullptr});
#else
                        attr_ptr[ip] = breit_wheeler_get_opt(amrex::RandomEngine{});
#endif
                    }
                }
            }
#endif
 
            for (int ia = 0; ia < n_user_real_attribs; ++ia)
            {
                // Current runtime comp is ia-th user defined attribute
                auto const it_ura = std::find(particle_comps.begin(), particle_comps.end(), user_real_attribs[ia]);
                if (it_ura != particle_comps.end() &&
                    std::distance(particle_comps.begin(), it_ura) == j)
                {
                    const amrex::ParserExecutor<7> user_real_attrib_parserexec =
                                             user_real_attrib_parser[ia]->compile<7>();
                    // If the particle tile was allocated in a memory pool that can run on GPU, launch GPU kernel
                    if (run_on_gpu) {
                        amrex::ParallelFor(stop - start,
                                           [=] AMREX_GPU_DEVICE (int i) noexcept {
                                               const int ip = i + start;
                                               amrex::ParticleReal xp, yp, zp;
                                               get_position(ip, xp, yp, zp);
                                               attr_ptr[ip] = user_real_attrib_parserexec(xp, yp, zp,
                                                                                          ux[ip], uy[ip], uz[ip], t);
                                           });
                    // Otherwise (e.g. particle tile allocated in pinned memory), run on CPU
                    } else {
                        for (int ip = start; ip < stop; ++ip) {
                            amrex::ParticleReal xp, yp, zp;
                            get_position(ip, xp, yp, zp);
                            attr_ptr[ip] = user_real_attrib_parserexec(xp, yp, zp,
                                                                       ux[ip], uy[ip], uz[ip], t);
                        }
                    }
                }
            }
        }
 
        // Initialize the last NumRuntimeIntComps() - n_external_attr_int runtime int attributes
        for (int j = n_external_attr_int; j < ptile.NumIntComps() ; ++j)
        {
            auto attr_ptr = ptile.GetStructOfArrays().GetIntData(j).data();
 
            // Current runtime comp is ionization level
            auto const it_ioniz = std::find(particle_icomps.begin(), particle_icomps.end(), "ionizationLevel");
            if (it_ioniz != particle_icomps.end() &&
                std::distance(particle_icomps.begin(), it_ioniz) == j)
            {
                if (run_on_gpu) {
                        amrex::ParallelFor(stop - start,
                                           [=] AMREX_GPU_DEVICE (int i) noexcept {
                                               const int ip = i + start;
                                               attr_ptr[ip] = ionization_initial_level;
                                           });
                } else {
                    for (int ip = start; ip < stop; ++ip) {
                        attr_ptr[ip] = ionization_initial_level;
                    }
                }
            }
 
            for (int ia = 0; ia < n_user_int_attribs; ++ia)
            {
                // Current runtime comp is ia-th user defined attribute
                auto const it_uia = std::find(particle_icomps.begin(), particle_icomps.end(), user_int_attribs[ia]);
                if (it_uia != particle_icomps.end() &&
                    std::distance(particle_icomps.begin(), it_uia) == j)
                {
                     const amrex::ParserExecutor<7> user_int_attrib_parserexec =
                                             user_int_attrib_parser[ia]->compile<7>();
                if (run_on_gpu) {
                        amrex::ParallelFor(stop - start,
                                           [=] AMREX_GPU_DEVICE (int i) noexcept {
                                               const int ip = i + start;
                                               amrex::ParticleReal xp, yp, zp;
                                               get_position(ip, xp, yp, zp);
                                               attr_ptr[ip] = static_cast<int>(
                                                                               user_int_attrib_parserexec(xp, yp, zp, ux[ip], uy[ip], uz[ip], t));
                                           });
                } else {
                    for (int ip = start; ip < stop; ++ip) {
                        amrex::ParticleReal xp, yp, zp;
                        get_position(ip, xp, yp, zp);
                        attr_ptr[ip] = static_cast<int>(
                                user_int_attrib_parserexec(xp, yp, zp, ux[ip], uy[ip], uz[ip], t));
                    }
                }
                }
            }
        }
}
 
}
 
#endif //WARPX_DEFAULTINITIALIZATION_H_