Update DSMC reaction weight calculation

Examples/Physics_applications/capacitive_discharge/PICMI_inputs_1d.py

@@ -12,6 +12,7 @@
 from scipy.sparse import linalg as sla
 
 from pywarpx import callbacks, fields, libwarpx, particle_containers, picmi
+from pywarpx.LoadThirdParty import load_cupy
 
 constants = picmi.constants
 
@@ -396,12 +397,14 @@ def rethermalize_neutrals(self):
         uy_arrays = self.neutral_cont.uyp
         uz_arrays = self.neutral_cont.uzp
 
+        xp, _ = load_cupy()
+
         vel_std = np.sqrt(constants.kb * self.gas_temp / self.m_ion)
         for ii in range(len(ux_arrays)):
             nps = len(ux_arrays[ii])
-            ux_arrays[ii][:] = vel_std * self.rng.normal(size=nps)
-            uy_arrays[ii][:] = vel_std * self.rng.normal(size=nps)
-            uz_arrays[ii][:] = vel_std * self.rng.normal(size=nps)
+            ux_arrays[ii][:] = xp.array(vel_std * self.rng.normal(size=nps))
+            uy_arrays[ii][:] = xp.array(vel_std * self.rng.normal(size=nps))
+            uz_arrays[ii][:] = xp.array(vel_std * self.rng.normal(size=nps))
 
     def _get_rho_ions(self):
         # deposit the ion density in rho_fp

Source/Particles/Collision/BinaryCollision/DSMC/DSMCFunc.H

@@ -110,11 +110,13 @@ public:
             amrex::ParticleReal * const AMREX_RESTRICT u1x = soa_1.m_rdata[PIdx::ux];
             amrex::ParticleReal * const AMREX_RESTRICT u1y = soa_1.m_rdata[PIdx::uy];
             amrex::ParticleReal * const AMREX_RESTRICT u1z = soa_1.m_rdata[PIdx::uz];
+            uint64_t* AMREX_RESTRICT idcpu1 = soa_1.m_idcpu;
 
             amrex::ParticleReal * const AMREX_RESTRICT w2 = soa_2.m_rdata[PIdx::w];
             amrex::ParticleReal * const AMREX_RESTRICT u2x = soa_2.m_rdata[PIdx::ux];
             amrex::ParticleReal * const AMREX_RESTRICT u2y = soa_2.m_rdata[PIdx::uy];
             amrex::ParticleReal * const AMREX_RESTRICT u2z = soa_2.m_rdata[PIdx::uz];
+            uint64_t* AMREX_RESTRICT idcpu2 = soa_2.m_idcpu;
 
             // Number of macroparticles of each species
             const index_type NI1 = I1e - I1s;
@@ -124,14 +126,9 @@ public:
 
             index_type pair_index = cell_start_pair + coll_idx;
 
-            // Because the number of particles of each species is not always equal (NI1 != NI2
-            // in general), some macroparticles will be paired with multiple macroparticles of the
-            // other species and we need to decrease their weight accordingly.
-            // c1 corresponds to the minimum number of times a particle of species 1 will be paired
-            // with a particle of species 2. Same for c2.
-            // index_type(1): https://github.com/AMReX-Codes/amrex/pull/3684
-            const index_type c1 = amrex::max(NI2/NI1, index_type(1));
-            const index_type c2 = amrex::max(NI1/NI2, index_type(1));
+            // multiplier ratio to take into account unsampled pairs
+            const auto multiplier_ratio = static_cast<int>(
+                m_isSameSpecies ? min_N + max_N - 1 : min_N);
 
 #if (defined WARPX_DIM_RZ)
             amrex::ParticleReal * const AMREX_RESTRICT theta1 = soa_1.m_rdata[PIdx::theta];
@@ -143,56 +140,64 @@ public:
             // stride (smaller set size) until we do all collisions (larger set size)
             for (index_type k = coll_idx; k < max_N; k += min_N)
             {
-                // c1k : how many times the current particle of species 1 is paired with a particle
-                // of species 2. Same for c2k.
-                const index_type c1k = (k%NI1 < max_N%NI1) ? c1 + 1: c1;
-                const index_type c2k = (k%NI2 < max_N%NI2) ? c2 + 1: c2;
+                // do not check for collision if a particle's weight was
+                // reduced to zero from a previous collision
+                if (idcpu1[ I1[i1] ]==amrex::ParticleIdCpus::Invalid ||
+                    idcpu2[ I2[i2] ]==amrex::ParticleIdCpus::Invalid ) {
+                    p_mask[pair_index] = false;
+                } else {
 
 #if (defined WARPX_DIM_RZ)
-                /* In RZ geometry, macroparticles can collide with other macroparticles
-                * in the same *cylindrical* cell. For this reason, collisions between macroparticles
-                * are actually not local in space. In this case, the underlying assumption is that
-                * particles within the same cylindrical cell represent a cylindrically-symmetry
-                * momentum distribution function. Therefore, here, we temporarily rotate the
-                * momentum of one of the macroparticles in agreement with this cylindrical symmetry.
-                * (This is technically only valid if we use only the m=0 azimuthal mode in the simulation;
-                * there is a corresponding assert statement at initialization.)
-                */
-                amrex::ParticleReal const theta = theta2[I2[i2]]-theta1[I1[i1]];
-                amrex::ParticleReal const u1xbuf = u1x[I1[i1]];
-                u1x[I1[i1]] = u1xbuf*std::cos(theta) - u1y[I1[i1]]*std::sin(theta);
-                u1y[I1[i1]] = u1xbuf*std::sin(theta) + u1y[I1[i1]]*std::cos(theta);
+                    /* In RZ geometry, macroparticles can collide with other macroparticles
+                    * in the same *cylindrical* cell. For this reason, collisions between macroparticles
+                    * are actually not local in space. In this case, the underlying assumption is that
+                    * particles within the same cylindrical cell represent a cylindrically-symmetry
+                    * momentum distribution function. Therefore, here, we temporarily rotate the
+                    * momentum of one of the macroparticles in agreement with this cylindrical symmetry.
+                    * (This is technically only valid if we use only the m=0 azimuthal mode in the simulation;
+                    * there is a corresponding assert statement at initialization.)
+                    */
+                    amrex::ParticleReal const theta = theta2[I2[i2]]-theta1[I1[i1]];
+                    amrex::ParticleReal const u1xbuf = u1x[I1[i1]];
+                    u1x[I1[i1]] = u1xbuf*std::cos(theta) - u1y[I1[i1]]*std::sin(theta);
+                    u1y[I1[i1]] = u1xbuf*std::sin(theta) + u1y[I1[i1]]*std::cos(theta);
 #endif
 
-                CollisionPairFilter(
-                    u1x[ I1[i1] ], u1y[ I1[i1] ], u1z[ I1[i1] ],
-                    u2x[ I2[i2] ], u2y[ I2[i2] ], u2z[ I2[i2] ],
-                    m1, m2, w1[ I1[i1] ]/c1k, w2[ I2[i2] ]/c2k,
-                    dt, dV, static_cast<int>(pair_index), p_mask,
-                    p_pair_reaction_weight, static_cast<int>(max_N),
-                    m_process_count, m_scattering_processes_data, engine);
+                    CollisionPairFilter(
+                        u1x[ I1[i1] ], u1y[ I1[i1] ], u1z[ I1[i1] ],
+                        u2x[ I2[i2] ], u2y[ I2[i2] ], u2z[ I2[i2] ],
+                        m1, m2, w1[ I1[i1] ], w2[ I2[i2] ],
+                        dt, dV, static_cast<int>(pair_index), p_mask,
+                        p_pair_reaction_weight, multiplier_ratio,
+                        m_process_count, m_scattering_processes_data, engine);
 
 #if (defined WARPX_DIM_RZ)
-                amrex::ParticleReal const u1xbuf_new = u1x[I1[i1]];
-                u1x[I1[i1]] = u1xbuf_new*std::cos(-theta) - u1y[I1[i1]]*std::sin(-theta);
-                u1y[I1[i1]] = u1xbuf_new*std::sin(-theta) + u1y[I1[i1]]*std::cos(-theta);
+                    amrex::ParticleReal const u1xbuf_new = u1x[I1[i1]];
+                    u1x[I1[i1]] = u1xbuf_new*std::cos(-theta) - u1y[I1[i1]]*std::sin(-theta);
+                    u1y[I1[i1]] = u1xbuf_new*std::sin(-theta) + u1y[I1[i1]]*std::cos(-theta);
 #endif
 
+                    // Remove pair reaction weight from the colliding particles' weights
+                    if (p_mask[pair_index]) {
+                        BinaryCollisionUtils::remove_weight_from_colliding_particle(
+                            w1[ I1[i1] ], idcpu1[ I1[i1] ], p_pair_reaction_weight[pair_index]);
+                        BinaryCollisionUtils::remove_weight_from_colliding_particle(
+                            w2[ I2[i2] ], idcpu2[ I2[i2] ], p_pair_reaction_weight[pair_index]);
+                    }
+                }
+
                 p_pair_indices_1[pair_index] = I1[i1];
                 p_pair_indices_2[pair_index] = I2[i2];
-                if (max_N == NI1) {
-                    i1 += min_N;
-                }
-                if (max_N == NI2) {
-                    i2 += min_N;
-                }
+                if (max_N == NI1) { i1 += min_N; }
+                if (max_N == NI2) { i2 += min_N; }
                 pair_index += min_N;
             }
         }
 
         int m_process_count;
         bool m_computeSpeciesDensities = false;
         bool m_computeSpeciesTemperatures = false;
+        bool m_isSameSpecies = false;
         ScatteringProcess::Executor* m_scattering_processes_data;
     };
 
@@ -201,6 +206,7 @@ public:
 private:
     amrex::Vector<ScatteringProcess> m_scattering_processes;
     amrex::Gpu::DeviceVector<ScatteringProcess::Executor> m_scattering_processes_exe;
+    bool m_isSameSpecies;
 
     Executor m_exe;
 };

Source/Particles/Collision/BinaryCollision/DSMC/DSMCFunc.cpp

@@ -19,7 +19,7 @@
 DSMCFunc::DSMCFunc (
     const std::string& collision_name,
     [[maybe_unused]] MultiParticleContainer const * const mypc,
-    [[maybe_unused]] const bool isSameSpecies )
+    const bool isSameSpecies ): m_isSameSpecies{isSameSpecies}
 {
     using namespace amrex::literals;
 
@@ -76,4 +76,5 @@ DSMCFunc::DSMCFunc (
     // Link executor to appropriate ScatteringProcess executors
     m_exe.m_scattering_processes_data = m_scattering_processes_exe.data();
     m_exe.m_process_count = process_count;
+    m_exe.m_isSameSpecies = m_isSameSpecies;
 }

Source/Particles/Collision/BinaryCollision/DSMC/SplitAndScatterFunc.H

@@ -94,11 +94,6 @@ public:
         const auto soa_1 = ptile1.getParticleTileData();
         const auto soa_2 = ptile2.getParticleTileData();
 
-        amrex::ParticleReal* AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
-        amrex::ParticleReal* AMREX_RESTRICT w2 = soa_2.m_rdata[PIdx::w];
-        uint64_t* AMREX_RESTRICT idcpu1 = soa_1.m_idcpu;
-        uint64_t* AMREX_RESTRICT idcpu2 = soa_2.m_idcpu;
-
         // Create necessary GPU vectors, that will be used in the kernel below
         amrex::Vector<SoaData_type> soa_products;
         for (int i = 0; i < m_num_product_species; i++)
@@ -151,12 +146,6 @@ public:
                 // Set the weight of the new particles to p_pair_reaction_weight[i]
                 soa_products_data[1].m_rdata[PIdx::w][product2_index] = p_pair_reaction_weight[i];
 
-                // Remove p_pair_reaction_weight[i] from the colliding particles' weights
-                BinaryCollisionUtils::remove_weight_from_colliding_particle(
-                    w1[p_pair_indices_1[i]], idcpu1[p_pair_indices_1[i]], p_pair_reaction_weight[i]);
-                BinaryCollisionUtils::remove_weight_from_colliding_particle(
-                    w2[p_pair_indices_2[i]], idcpu2[p_pair_indices_2[i]], p_pair_reaction_weight[i]);
-
                 // Set the child particle properties appropriately
                 auto& ux1 = soa_products_data[0].m_rdata[PIdx::ux][product1_index];
                 auto& uy1 = soa_products_data[0].m_rdata[PIdx::uy][product1_index];