Update DSMC reaction weight calculation (ECP-WarpX#5135)

* initial changes to reaction weight calculation following ECP-WarpX#5133 * remove weight reduction logic in `SplitAndScatterFunc.H` Signed-off-by: roelof-groenewald <regroenewald@gmail.com> * update 1d DSMC example to be runable on NVIDIA GPU (uses `cupy`) * add `m_isSameSpecies` member variable to `DSMCFunc`; reduce code duplication * reset DSMC test checksum values --------- Signed-off-by: roelof-groenewald <regroenewald@gmail.com>
JustinRayAngus · Aug 19, 2024 · f02aaa2 · f02aaa2
1 parent 6007f8a
commit f02aaa2
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 107 deletions.
diff --git a/Examples/Physics_applications/capacitive_discharge/PICMI_inputs_1d.py b/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

diff --git a/Examples/Physics_applications/capacitive_discharge/analysis_dsmc.py b/Examples/Physics_applications/capacitive_discharge/analysis_dsmc.py
@@ -18,39 +18,39 @@
 my_check = checksumAPI.evaluate_checksum(test_name, fn, do_particles=True)
 
 ref_density = np.array([
-    1.27943881e+14, 2.23583097e+14, 2.55396716e+14, 2.55673406e+14,
-    2.55827566e+14, 2.55803446e+14, 2.55798707e+14, 2.55748961e+14,
-    2.55906413e+14, 2.56063991e+14, 2.55937018e+14, 2.55841390e+14,
-    2.55917724e+14, 2.55988641e+14, 2.56052050e+14, 2.56285151e+14,
-    2.56647960e+14, 2.56756264e+14, 2.56430158e+14, 2.56117493e+14,
-    2.56065302e+14, 2.56265220e+14, 2.56328575e+14, 2.56031495e+14,
-    2.56123757e+14, 2.56431173e+14, 2.56385320e+14, 2.56391170e+14,
-    2.56561177e+14, 2.56513926e+14, 2.56332201e+14, 2.56252442e+14,
-    2.56238982e+14, 2.56216498e+14, 2.56461281e+14, 2.56863199e+14,
-    2.56908100e+14, 2.56926112e+14, 2.57001641e+14, 2.56735963e+14,
-    2.56315358e+14, 2.56137028e+14, 2.56101418e+14, 2.56276827e+14,
-    2.56425668e+14, 2.56181798e+14, 2.56044925e+14, 2.56330387e+14,
-    2.56623150e+14, 2.56445316e+14, 2.56292750e+14, 2.56440918e+14,
-    2.56433406e+14, 2.56186982e+14, 2.56236390e+14, 2.56469557e+14,
-    2.56349704e+14, 2.56487457e+14, 2.56771823e+14, 2.56614683e+14,
-    2.56552210e+14, 2.56850291e+14, 2.56783396e+14, 2.56483187e+14,
-    2.56510868e+14, 2.56490408e+14, 2.56656042e+14, 2.56820924e+14,
-    2.56640314e+14, 2.56465063e+14, 2.56510264e+14, 2.56917331e+14,
-    2.57228490e+14, 2.56960593e+14, 2.56587911e+14, 2.56672682e+14,
-    2.56774414e+14, 2.56548335e+14, 2.56225540e+14, 2.56079693e+14,
-    2.56062796e+14, 2.56054612e+14, 2.56028683e+14, 2.56068820e+14,
-    2.56380975e+14, 2.56654914e+14, 2.56776792e+14, 2.56983661e+14,
-    2.56989477e+14, 2.56646250e+14, 2.56589639e+14, 2.56946205e+14,
-    2.57091201e+14, 2.56913590e+14, 2.56513535e+14, 2.56122597e+14,
-    2.56176340e+14, 2.56808001e+14, 2.57239393e+14, 2.56845066e+14,
-    2.56662482e+14, 2.56862583e+14, 2.56518922e+14, 2.56155531e+14,
-    2.56362794e+14, 2.57203564e+14, 2.57737938e+14, 2.57252026e+14,
-    2.56859277e+14, 2.56658995e+14, 2.56357364e+14, 2.56393454e+14,
-    2.56714308e+14, 2.57042200e+14, 2.57551087e+14, 2.57502490e+14,
-    2.56641118e+14, 2.56401115e+14, 2.56644629e+14, 2.56673096e+14,
-    2.56534659e+14, 2.56357745e+14, 2.56455309e+14, 2.56586850e+14,
-    2.56442415e+14, 2.56335971e+14, 2.56411429e+14, 2.24109018e+14,
-    1.27678869e+14
+    1.27942709e+14, 2.23579371e+14, 2.55384387e+14, 2.55660663e+14,
+    2.55830911e+14, 2.55814337e+14, 2.55798906e+14, 2.55744891e+14,
+    2.55915585e+14, 2.56083194e+14, 2.55942354e+14, 2.55833026e+14,
+    2.56036175e+14, 2.56234141e+14, 2.56196179e+14, 2.56146141e+14,
+    2.56168022e+14, 2.56216909e+14, 2.56119961e+14, 2.56065167e+14,
+    2.56194764e+14, 2.56416398e+14, 2.56465239e+14, 2.56234337e+14,
+    2.56234503e+14, 2.56316003e+14, 2.56175023e+14, 2.56030269e+14,
+    2.56189301e+14, 2.56286379e+14, 2.56130396e+14, 2.56295225e+14,
+    2.56474082e+14, 2.56340375e+14, 2.56350864e+14, 2.56462330e+14,
+    2.56469391e+14, 2.56412726e+14, 2.56241788e+14, 2.56355650e+14,
+    2.56650599e+14, 2.56674748e+14, 2.56642480e+14, 2.56823508e+14,
+    2.57025029e+14, 2.57110614e+14, 2.57042364e+14, 2.56950884e+14,
+    2.57051822e+14, 2.56952148e+14, 2.56684016e+14, 2.56481130e+14,
+    2.56277073e+14, 2.56065774e+14, 2.56190033e+14, 2.56411074e+14,
+    2.56202418e+14, 2.56128368e+14, 2.56227002e+14, 2.56083004e+14,
+    2.56056768e+14, 2.56343831e+14, 2.56443659e+14, 2.56280541e+14,
+    2.56191572e+14, 2.56147304e+14, 2.56342794e+14, 2.56735473e+14,
+    2.56994680e+14, 2.56901500e+14, 2.56527131e+14, 2.56490824e+14,
+    2.56614730e+14, 2.56382744e+14, 2.56588214e+14, 2.57160270e+14,
+    2.57230435e+14, 2.57116530e+14, 2.57065771e+14, 2.57236507e+14,
+    2.57112865e+14, 2.56540177e+14, 2.56416828e+14, 2.56648954e+14,
+    2.56625594e+14, 2.56411003e+14, 2.56523754e+14, 2.56841108e+14,
+    2.56856368e+14, 2.56757912e+14, 2.56895134e+14, 2.57144419e+14,
+    2.57001944e+14, 2.56371759e+14, 2.56179404e+14, 2.56541905e+14,
+    2.56715727e+14, 2.56851681e+14, 2.57114458e+14, 2.57001739e+14,
+    2.56825690e+14, 2.56879682e+14, 2.56699673e+14, 2.56532841e+14,
+    2.56479582e+14, 2.56630989e+14, 2.56885996e+14, 2.56694637e+14,
+    2.56250819e+14, 2.56045278e+14, 2.56366075e+14, 2.56693733e+14,
+    2.56618530e+14, 2.56580918e+14, 2.56812781e+14, 2.56754216e+14,
+    2.56444736e+14, 2.56473391e+14, 2.56538398e+14, 2.56626551e+14,
+    2.56471950e+14, 2.56274969e+14, 2.56489423e+14, 2.56645266e+14,
+    2.56611124e+14, 2.56344324e+14, 2.56244156e+14, 2.24183727e+14,
+    1.27909856e+14
 ])
 
 density_data = np.load( 'ion_density_case_1.npy' )

diff --git a/Regression/Checksum/benchmarks_json/Python_dsmc_1d.json b/Regression/Checksum/benchmarks_json/Python_dsmc_1d.json
@@ -1,27 +1,27 @@
 {
   "lev=0": {
-    "rho_electrons": 0.004434311371276535,
-    "rho_he_ions": 0.005200489146365628
+    "rho_electrons": 0.004437338851654305,
+    "rho_he_ions": 0.005200276265886133
   },
   "he_ions": {
-    "particle_momentum_x": 2.7688639005236794e-19,
-    "particle_momentum_y": 2.760648864116014e-19,
-    "particle_momentum_z": 3.6226628630061563e-19,
-    "particle_position_x": 2201.614485497906,
-    "particle_weight": 17190996093750.002
+    "particle_momentum_x": 2.768463746716725e-19,
+    "particle_momentum_y": 2.7585450668167785e-19,
+    "particle_momentum_z": 3.6189671443598533e-19,
+    "particle_position_x": 2201.408357891233,
+    "particle_weight": 17190472656250.002
+  },
+  "electrons": {
+    "particle_momentum_x": 3.523554668287801e-20,
+    "particle_momentum_y": 3.515628626179393e-20,
+    "particle_momentum_z": 1.258711379033217e-19,
+    "particle_position_x": 2140.8168584833174,
+    "particle_weight": 14588988281250.002
   },
   "neutrals": {
-    "particle_momentum_x": 1.4040775167811838e-19,
-    "particle_momentum_y": 1.4009514702703257e-19,
-    "particle_momentum_z": 1.4093144247152345e-19,
-    "particle_position_x": 1119.524782452131,
+    "particle_momentum_x": 1.4054952479597137e-19,
+    "particle_momentum_y": 1.403311018061206e-19,
+    "particle_momentum_z": 1.411491089895956e-19,
+    "particle_position_x": 1119.82858839282,
     "particle_weight": 6.4588e+19
-  },
-  "electrons": {
-    "particle_momentum_x": 3.5193566020597954e-20,
-    "particle_momentum_y": 3.5368803263788353e-20,
-    "particle_momentum_z": 1.2572273121326582e-19,
-    "particle_position_x": 2139.3709122461873,
-    "particle_weight": 14579566406250.002
   }
 }
diff --git a/Source/Particles/Collision/BinaryCollision/DSMC/DSMCFunc.H b/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;
 };

diff --git a/Source/Particles/Collision/BinaryCollision/DSMC/DSMCFunc.cpp b/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;
 }
diff --git a/Source/Particles/Collision/BinaryCollision/DSMC/SplitAndScatterFunc.H b/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];