random-batch-ewald

Author: brucefan1983

src/force/nep_charge.cu

@@ -27,6 +27,7 @@ heat transport, Phys. Rev. B. 104, 104309 (2021).
 #include "utilities/error.cuh"
 #include "utilities/gpu_macro.cuh"
 #include "utilities/nep_utilities.cuh"
+#include <chrono>
 #include <cmath>
 #include <cstring>
 #include <fstream>
@@ -74,6 +75,12 @@ void NEP_Charge::initialize_dftd3()
 
 NEP_Charge::NEP_Charge(const char* file_potential, const int num_atoms)
 {
+  #ifdef DEBUG
+  rng = std::mt19937(12345678);
+#else
+  rng = std::mt19937(std::chrono::system_clock::now().time_since_epoch().count());
+#endif
+
   std::ifstream input(file_potential);
   if (!input.is_open()) {
     std::cout << "Failed to open " << file_potential << std::endl;
@@ -1259,6 +1266,40 @@ void NEP_Charge::find_k_and_G(const double* box)
   std::vector<float> cpu_kz;
   std::vector<float> cpu_G;
 
+#ifdef USE_RBE
+  std::uniform_real_distribution<float> rand_number(0.0f, 1.0f);
+  float normalization_factor = 0.0f;
+  for (int n1 = 0; n1 <= n1_max; ++n1) {
+    for (int n2 = - n2_max; n2 <= n2_max; ++n2) {
+      for (int n3 = - n3_max; n3 <= n3_max; ++n3) {
+        const int nsq = n1 * n1 + n2 * n2 + n3 * n3;
+        if (nsq == 0 || (n1 == 0 && n2 < 0) || (n1 == 0 && n2 == 0 && n3 < 0)) continue;
+        const float kx = n1 * b1[0] + n2 * b2[0] + n3 * b3[0];
+        const float ky = n1 * b1[1] + n2 * b2[1] + n3 * b3[1];
+        const float kz = n1 * b1[2] + n2 * b2[2] + n3 * b3[2];
+        const float ksq = kx * kx + ky * ky + kz * kz;
+        if (ksq < ksq_max) {
+          float exp_factor = exp(-ksq * charge_para.alpha_factor);
+          normalization_factor += exp_factor;
+          if (rand_number(rng) < exp_factor) {
+            cpu_kx.emplace_back(kx);
+            cpu_ky.emplace_back(ky);
+            cpu_kz.emplace_back(kz);
+            float G = abs(two_pi_over_det) / ksq;
+            cpu_G.emplace_back(2.0f * G);
+          }
+        }
+      }
+    }
+  }
+
+  int num_kpoints = int(cpu_kx.size());
+  for (int n = 0; n < num_kpoints; ++n) {
+    cpu_G[n] *= normalization_factor / num_kpoints;
+  }
+
+#else
+
   for (int n1 = 0; n1 <= n1_max; ++n1) {
     for (int n2 = - n2_max; n2 <= n2_max; ++n2) {
       for (int n3 = - n3_max; n3 <= n3_max; ++n3) {
@@ -1280,6 +1321,8 @@ void NEP_Charge::find_k_and_G(const double* box)
   }
 
   int num_kpoints = int(cpu_kx.size());
+#endif
+
   if (num_kpoints > charge_para.num_kpoints_max) {
     charge_para.num_kpoints_max = num_kpoints;
     nep_data.kx.resize(charge_para.num_kpoints_max);

src/force/nep_charge.cuh

@@ -18,6 +18,7 @@
 #include "potential.cuh"
 #include "utilities/common.cuh"
 #include "utilities/gpu_vector.cuh"
+#include <random>
 
 struct NEP_Charge_Data {
   GPU_Vector<float> f12x; // 3-body or manybody partial forces
@@ -146,6 +147,7 @@ private:
   ExpandedBox ebox;
   DFTD3 dftd3;
   Charge_Para charge_para;
+  std::mt19937 rng;
 
   void update_potential(float* parameters, ANN& ann);
 #ifdef USE_TABLE