Merge branch 'dev' into comm

4 files changed, 397 insertions, 0 deletions

diff --git a/Source/Particles/ParticleCreation/CopyParticle.H b/Source/Particles/ParticleCreation/CopyParticle.H
new file mode 100644
index 000000000..bd2d530fb
--- /dev/null
+++ b/Source/Particles/ParticleCreation/CopyParticle.H
@@ -0,0 +1,90 @@
+#ifndef COPYPARTICLE_H_
+#define COPYPARTICLE_H_
+
+#include "WarpXParticleContainer.H"
+
+/**
+ * \brief Functor to copy one particle
+ *
+ * This is meant to be a very small class captured by value in kernel launches,
+ * that can be initialized on the host and copied to the device at each
+ * iteration. It should be general enough to be used by all processes.
+ *
+ * Pointers to SoA data are saved when constructor is called.
+ * AoS data is passed as argument of operator ().
+ */
+class copyParticle
+{
+
+public:
+
+    // ID of MPI rank
+    int m_cpuid;
+    // If true, will copy old attribs for back-transforme diagnostics
+    bool m_do_back_transformed_product;
+    // Source old (runtime) attribs for back-transformed diagnostics
+    amrex::GpuArray<amrex::ParticleReal*,3> m_runtime_uold_source;
+    // Source attribs
+    amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> m_attribs_source;
+    // Product attribs
+    amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> m_attribs_product;
+    // Product runtime attribs
+    amrex::GpuArray<amrex::ParticleReal*,6> m_runtime_attribs_product;
+
+    // Simple constructor
+    AMREX_GPU_HOST_DEVICE
+    copyParticle(
+        const int cpuid, const int do_back_transformed_product,
+        const amrex::GpuArray<amrex::ParticleReal*,3> runtime_uold_source,
+        const amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_source,
+        const amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_product,
+        const amrex::GpuArray<amrex::ParticleReal*,6> runtime_attribs_product)
+        :
+        m_cpuid(cpuid),
+        m_do_back_transformed_product(do_back_transformed_product),
+        m_runtime_uold_source(runtime_uold_source),
+        m_attribs_source(attribs_source),
+        m_attribs_product(attribs_product),
+        m_runtime_attribs_product(runtime_attribs_product){}
+
+    /**
+     * \brief Overload operator () to do the copy of 1 particle
+     *
+     * \param is index of source particle
+     * \param ip index of product particle
+     * \param pid_product ID of first product particle
+     * \param p_source Struct with data for 1 source particle (position etc.)
+     * \param p_source Struct with data for 1 product particle (position etc.)
+     */
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    void operator () (int is, int ip, int pid_product,
+                      WarpXParticleContainer::ParticleType& p_source,
+                      WarpXParticleContainer::ParticleType& p_product) const noexcept
+    {
+        // Copy particle from source to product: AoS
+        p_product.id() = pid_product + ip;
+        p_product.cpu() = m_cpuid;
+        p_product.pos(0) = p_source.pos(0);
+        p_product.pos(1) = p_source.pos(1);
+#if (AMREX_SPACEDIM == 3)
+        p_product.pos(2) = p_source.pos(2);
+#endif
+        // Copy particle from source to product: SoA
+        for (int ia = 0; ia < PIdx::nattribs; ++ia) {
+            m_attribs_product[ia][ip] = m_attribs_source[ia][is];
+        }
+        // Update xold etc. if boosted frame diagnostics required
+        // for product species. Fill runtime attribs with a copy of
+        // current properties (xold = x etc.).
+        if (m_do_back_transformed_product) {
+            m_runtime_attribs_product[0][ip] = p_source.pos(0);
+            m_runtime_attribs_product[1][ip] = p_source.pos(1);
+            m_runtime_attribs_product[2][ip] = p_source.pos(2);
+            m_runtime_attribs_product[3][ip] = m_runtime_uold_source[0][ip];
+            m_runtime_attribs_product[4][ip] = m_runtime_uold_source[1][ip];
+            m_runtime_attribs_product[5][ip] = m_runtime_uold_source[2][ip];
+        }
+    }
+};
+
+#endif // COPYPARTICLE_H_
diff --git a/Source/Particles/ParticleCreation/ElementaryProcess.H b/Source/Particles/ParticleCreation/ElementaryProcess.H
new file mode 100644
index 000000000..fa791c244
--- /dev/null
+++ b/Source/Particles/ParticleCreation/ElementaryProcess.H
@@ -0,0 +1,257 @@
+#ifndef ELEMENTARYPROCESS_H_
+#define ELEMENTARYPROCESS_H_
+
+#include "WarpXParticleContainer.H"
+#include "CopyParticle.H"
+#include "TransformParticle.H"
+
+/**
+ * \brief Base class for particle creation processes
+ *
+ * Particles in a source particle container are copied to product particle
+ * containers if they are flagged. Both source and product particles can be
+ * modified.
+ *
+ * - initialize_copy initializes a general copy functor
+ * - createParticles formats the data to prepare for the copy, then
+ *       calls copyAndTransformParticles
+ * - copyAndTransformParticles loops over particles, performs the copy and
+ *       transform source and product particles if needed.
+ *
+ * The class is templated on the process type. This gives flexibility
+ * on source and product particle transformations.
+ */
+template<elementaryProcessType ProcessT>
+class elementaryProcess
+{
+public:
+
+    /**
+     * \brief initialize and return functor for particle copy
+     *
+     * \param cpuid id of MPI rank
+     * \param do_back_transformed_product whether to copy old attribs
+     * \param runtime_uold_source momentum of source particles
+     * \param attribs_source attribs of source particles
+     * \param attribs_product attribs of product particles
+     * \param runtime_attribs_product runtime attribs for product, to store old attribs
+     */
+    copyParticle initialize_copy(
+        const int cpuid, const int do_back_transformed_product,
+        const amrex::GpuArray<amrex::ParticleReal*,3> runtime_uold_source,
+        const amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_source,
+        const amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_product,
+        const amrex::GpuArray<amrex::ParticleReal*,6> runtime_attribs_product) const noexcept
+    {
+        return copyParticle (
+            cpuid, do_back_transformed_product, runtime_uold_source,
+            attribs_source, attribs_product, runtime_attribs_product);
+    };
+
+    /**
+     * \brief create particles in product particle containers
+     *
+     * For particle i in mfi, if is_flagged[i]=1, copy particle
+     * particle i from container pc_source into each container in v_pc_product
+     *
+     * \param lev MR level
+     * \param mfi MFIter where source particles are located
+     * \param pc_source source particle container
+     * \param v_pc_product vector of product particle containers
+     * \param is_flagged particles for which is_flagged is 1 are copied
+     * \param v_do_back_transformed_product vector: whether to copy old attribs for
+     *        each product container
+     */
+    void createParticles (
+        int lev, const amrex::MFIter& mfi,
+        std::unique_ptr< WarpXParticleContainer>& pc_source,
+        amrex::Vector<WarpXParticleContainer*> v_pc_product,
+        amrex::Gpu::ManagedDeviceVector<int>& is_flagged,
+        amrex::Gpu::ManagedDeviceVector<int> v_do_back_transformed_product)
+    {
+
+        BL_PROFILE("createParticles");
+        const int grid_id = mfi.index();
+        const int tile_id = mfi.LocalTileIndex();
+
+        // Get source particle data
+        auto& ptile_source = pc_source->GetParticles(lev)[std::make_pair(grid_id,tile_id)];
+        const int np_source = ptile_source.GetArrayOfStructs().size();
+        if (np_source == 0) return;
+        // --- source AoS particle data
+        WarpXParticleContainer::ParticleType* particles_source = ptile_source.GetArrayOfStructs()().data();
+        // --- source SoA particle data
+        auto& soa_source = ptile_source.GetStructOfArrays();
+        amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_source;
+        for (int ia = 0; ia < PIdx::nattribs; ++ia) {
+            attribs_source[ia] = soa_source.GetRealData(ia).data();
+        }
+        // --- source runtime attribs
+        amrex::GpuArray<amrex::ParticleReal*,3> runtime_uold_source;
+        // Prepare arrays for boosted frame diagnostics.
+        runtime_uold_source[0] = soa_source.GetRealData(PIdx::ux).data();
+        runtime_uold_source[1] = soa_source.GetRealData(PIdx::uy).data();
+        runtime_uold_source[2] = soa_source.GetRealData(PIdx::uz).data();
+
+        // --- source runtime integer attribs
+        amrex::GpuArray<int*,1> runtime_iattribs_source;
+        std::map<std::string, int> icomps_source = pc_source->getParticleiComps();
+        runtime_iattribs_source[0] = soa_source.GetIntData(icomps_source["ionization_level"]).data();
+
+        int nproducts = v_pc_product.size();
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+            v_do_back_transformed_product.size() == nproducts,
+            "v_pc_product and v_do_back_transformed_product must have the same size.");
+
+        // Indices of product particle for each source particle.
+        // i_product[i]-1 is the location in product tile of product particle
+        // from source particle i.
+        amrex::Gpu::ManagedDeviceVector<int> i_product;
+        i_product.resize(np_source);
+        // 0<i<np_source
+        // 0<i_product<np_flagged
+        // Strictly speaking, i_product should be an exclusive_scan of
+        // is_flagged. However, for indices where is_flagged is 1, the
+        // inclusive scan gives the same result with an offset of 1.
+        // The advantage of inclusive_scan is that the sum of is_flagged
+        // is in the last element, so no other reduction is required to get
+        // number of particles.
+        // Gpu::inclusive_scan runs on the current GPU stream, and synchronizes
+        // with the CPU, so that the next line (executed by the CPU) has the
+        // updated values of i_product
+        amrex::Gpu::inclusive_scan(is_flagged.begin(), is_flagged.end(), i_product.begin());
+        int np_flagged = i_product[np_source-1];
+        if (np_flagged == 0) return;
+
+        amrex::Vector<copyParticle> v_copy_functor;
+        amrex::Gpu::ManagedDeviceVector<int> v_pid_product(nproducts);
+        amrex::Gpu::ManagedDeviceVector<WarpXParticleContainer::ParticleType*> v_particles_product(nproducts);
+        for (int iproduct=0; iproduct<nproducts; iproduct++){
+            WarpXParticleContainer*& pc_product = v_pc_product[iproduct];
+            // Get product particle data
+            auto& ptile_product = pc_product->GetParticles(lev)[std::make_pair(grid_id,tile_id)];
+            // old and new (i.e., including flagged particles) number of particles
+            // for product species
+            const int np_product_old = ptile_product.GetArrayOfStructs().size();
+            const int np_product_new = np_product_old + np_flagged;
+            // Allocate extra space in product species for flagged particles.
+            ptile_product.resize(np_product_new);
+            // --- product AoS particle data
+            // WarpXParticleContainer::ParticleType* particles_product = ptile_product.GetArrayOfStructs()().data() + np_product_old;
+            v_particles_product[iproduct] = ptile_product.GetArrayOfStructs()().data() + np_product_old;
+            // First element is the first newly-created product particle
+            // --- product SoA particle data
+            auto& soa_product = ptile_product.GetStructOfArrays();
+            amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_product;
+            for (int ia = 0; ia < PIdx::nattribs; ++ia) {
+                // First element is the first newly-created product particle
+                attribs_product[ia] = soa_product.GetRealData(ia).data() + np_product_old;
+            }
+            // --- product runtime attribs
+            amrex::GpuArray<amrex::ParticleReal*,6> runtime_attribs_product;
+            const int do_boost = v_do_back_transformed_product[iproduct];
+            if (do_boost) {
+                std::map<std::string, int> comps_product = pc_product->getParticleComps();
+                runtime_attribs_product[0] = soa_product.GetRealData(comps_product[ "xold"]).data() + np_product_old;
+                runtime_attribs_product[1] = soa_product.GetRealData(comps_product[ "yold"]).data() + np_product_old;
+                runtime_attribs_product[2] = soa_product.GetRealData(comps_product[ "zold"]).data() + np_product_old;
+                runtime_attribs_product[3] = soa_product.GetRealData(comps_product["uxold"]).data() + np_product_old;
+                runtime_attribs_product[4] = soa_product.GetRealData(comps_product["uyold"]).data() + np_product_old;
+                runtime_attribs_product[5] = soa_product.GetRealData(comps_product["uzold"]).data() + np_product_old;
+            }
+
+            // --- product runtime integer attribs
+            int pid_product;
+#pragma omp critical (doParticleCreation_nextid)
+            {
+                // ID of first newly-created product particle
+                pid_product = pc_product->NextID();
+                // Update NextID to include particles created in this function
+                pc_product->setNextID(pid_product+np_flagged);
+            }
+            const int cpuid = amrex::ParallelDescriptor::MyProc();
+
+            // Create instance of copy functor, and add it to the vector
+            v_copy_functor.push_back (initialize_copy(
+                                          cpuid, v_do_back_transformed_product[iproduct],
+                                          runtime_uold_source,
+                                          attribs_source,
+                                          attribs_product,
+                                          runtime_attribs_product) );
+            v_pid_product[iproduct] = pid_product;
+
+        }
+        // Loop over source particles and create product particles
+        copyAndTransformParticles(is_flagged, i_product, np_source, v_pid_product,
+                                  v_particles_product, particles_source, v_copy_functor,
+                                  runtime_iattribs_source);
+        // Synchronize to prevent the destruction of temporary arrays (at the
+        // end of the function call) before the kernel executes.
+        amrex::Gpu::streamSynchronize();
+    }
+
+    /**
+     * \brief Loop over source particles and create product particles
+     *
+     * \param is_flagged particles for which is_flagged is 1 are copied
+     * \param i_product relative indices of product particle. This is the same
+     *        for all product particle containers.
+     * \param np_source number of particles in source container
+     * \param v_pid_product for each product particle container: ID of the
+     *        first product particle. Others IDs are incremented from this one
+     * \param v_particles_product for each product particle container: pointer
+     *        to the AoS particle data
+     * \param particles_source pointter to the AoS source particle data
+     * \param v_copy_functor vector of copy functors, one for each product particle container
+     * \param runtime_iattribs_source pointer to source runtime integer attribs
+     */
+    void copyAndTransformParticles(
+        amrex::Gpu::ManagedDeviceVector<int>& is_flagged,
+        amrex::Gpu::ManagedDeviceVector<int>& i_product,
+        int np_source,
+        amrex::Gpu::ManagedDeviceVector<int> v_pid_product,
+        amrex::Gpu::ManagedDeviceVector<WarpXParticleContainer::ParticleType*> v_particles_product,
+        WarpXParticleContainer::ParticleType* particles_source,
+        const amrex::Vector<copyParticle>& v_copy_functor,
+        amrex::GpuArray<int*,1> runtime_iattribs_source)
+    {
+        int const * const AMREX_RESTRICT p_is_flagged = is_flagged.dataPtr();
+        int const * const AMREX_RESTRICT p_i_product = i_product.dataPtr();
+        copyParticle const * const p_copy_functor = v_copy_functor.data();
+        WarpXParticleContainer::ParticleType * * p_particles_product = v_particles_product.data();
+        int const * const p_pid_product = v_pid_product.data();
+
+        // Loop over all source particles. If is_flagged, copy particle data
+        // to corresponding product particle.
+        int nproducts = v_particles_product.size();
+        amrex::For(
+            np_source, [=] AMREX_GPU_DEVICE (int is) noexcept
+            {
+                if(p_is_flagged[is]){
+                    // offset of 1 due to inclusive scan
+                    int ip = p_i_product[is]-1;
+                    WarpXParticleContainer::ParticleType& p_source  = particles_source[is];
+                    for (int iproduct=0; iproduct<nproducts; iproduct++){
+                        // is: index of flagged particle in source species
+                        // ip: index of corresponding new particle in product species
+                        WarpXParticleContainer::ParticleType* particles_product = p_particles_product[iproduct];
+                        WarpXParticleContainer::ParticleType& p_product = particles_product[ip];
+                        p_copy_functor[iproduct](is, ip, p_pid_product[iproduct], p_source, p_product);
+                    }
+                    // Transform source particle
+                    transformSourceParticle<ProcessT>(is, p_source, runtime_iattribs_source);
+                    // Transform product particles. The loop over product
+                    // species is done inside the function to allow for
+                    // more flexibility.
+                    transformProductParticle<ProcessT>(ip, p_particles_product);
+                }
+            }
+        );
+    }
+};
+
+// Derived class for ionization process
+class IonizationProcess: public elementaryProcess<elementaryProcessType::Ionization>
+{};
+
+#endif // ELEMENTARYPROCESS_H_
diff --git a/Source/Particles/ParticleCreation/Make.package b/Source/Particles/ParticleCreation/Make.package
new file mode 100644
index 000000000..6e32f4a77
--- /dev/null
+++ b/Source/Particles/ParticleCreation/Make.package
@@ -0,0 +1,6 @@
+CEXE_headers += ElementaryProcess.H
+CEXE_headers += CopyParticle.H
+CEXE_headers += TransformParticle.H
+
+INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/Particles/ParticleCreation/
+VPATH_LOCATIONS   += $(WARPX_HOME)/Source/Particles/ParticleCreation/
diff --git a/Source/Particles/ParticleCreation/TransformParticle.H b/Source/Particles/ParticleCreation/TransformParticle.H
new file mode 100644
index 000000000..14e534d0e
--- /dev/null
+++ b/Source/Particles/ParticleCreation/TransformParticle.H
@@ -0,0 +1,44 @@
+#ifndef TRANSFORMPARTICLE_H_
+#define TRANSFORMPARTICLE_H_
+
+#include "WarpXParticleContainer.H"
+
+enum struct elementaryProcessType { Ionization };
+
+/**
+ * \brief small modifications on source particle
+ * \param i index of particle
+ * \param particle particle struct
+ * \param runtime_iattribs integer attribs
+ */
+template < elementaryProcessType ProcessT >
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+void transformSourceParticle(
+    int i,
+    WarpXParticleContainer::ParticleType& particle,
+    amrex::GpuArray<int*,1> runtime_iattribs){}
+
+/**
+ * \brief small modifications on product particle
+ * \param i index of particle
+ * \param particle particle struct
+ */
+template < elementaryProcessType ProcessT >
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+void transformProductParticle(
+    int i,
+    WarpXParticleContainer::ParticleType* * v_particle){}
+
+// For ionization, increase ionization level of source particle
+template <>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+void transformSourceParticle < elementaryProcessType::Ionization > (
+    int i,
+    WarpXParticleContainer::ParticleType& particle,
+    amrex::GpuArray<int*,1> runtime_iattribs)
+{
+    // increment particle's ionization level
+    runtime_iattribs[0][i] += 1;
+}
+
+#endif // TRANSFORMPARTICLE_H_