Merge pull request brucefan1983#1039 from brucefan1983/ref_total_charge

use ref charge

Author: brucefan1983

src/main_nep/dataset.cu

@@ -897,7 +897,7 @@ static __global__ void gpu_sum_charge_error(
   }
 
   if (tid == 0) {
-    float diff = s_charge[0] / Na - g_charge_ref[bid];
+    float diff = (s_charge[0] - g_charge_ref[bid]) / Na;
     error_gpu[bid] = diff * diff;
   }
 }

src/main_nep/nep_charge.cu

@@ -1075,7 +1075,8 @@ static __global__ void find_k_and_G(
 static __global__ void zero_total_charge(
   const int* Na,
   const int* Na_sum,
-  float* g_charge,
+  const float* g_charge_ref,
+  const float* g_charge,
   float* g_charge_shifted)
 {
   int tid = threadIdx.x;
@@ -1103,7 +1104,7 @@ static __global__ void zero_total_charge(
   for (int batch = 0; batch < number_of_batches; ++batch) {
     int n = tid + batch * 1024 + N1;
     if (n < N2) {
-      g_charge_shifted[n] = g_charge[n] - s_charge[0] / (N2 - N1);
+      g_charge_shifted[n] = g_charge[n] + (g_charge_ref[blockIdx.x] - s_charge[0]) / (N2 - N1);
     }
   }
 }
@@ -1245,14 +1246,16 @@ void NEP_Charge::find_force(
       nep_data[device_id].charge_derivative.data());
     GPU_CHECK_KERNEL
 
-    // enforce charge neutrality
+    // enforce total charge is the target
     zero_total_charge<<<dataset[device_id].Nc, 1024>>>(
       dataset[device_id].Na.data(),
       dataset[device_id].Na_sum.data(),
+      dataset[device_id].charge_ref_gpu.data(),
       dataset[device_id].charge.data(),
       dataset[device_id].charge_shifted.data());
     GPU_CHECK_KERNEL
 
+    // modes 1 and 2 have reciprocal space
     if (paramb.charge_mode != 3) {
       find_k_and_G<<<(dataset[device_id].Nc - 1) / 64 + 1, 64>>>(
         dataset[device_id].Nc,
@@ -1309,8 +1312,7 @@ void NEP_Charge::find_force(
       GPU_CHECK_KERNEL
     }
 
-    // charge_mode = 1: include real space and self energy
-    // charge_mode = 2: exclude real space and self energy
+    // mode 1 has real space
     if (paramb.charge_mode == 1) {
       find_force_charge_real_space<<<grid_size, block_size>>>(
         dataset[device_id].N,
@@ -1329,7 +1331,10 @@ void NEP_Charge::find_force(
         dataset[device_id].energy.data(),
         nep_data[device_id].D_real.data());
       GPU_CHECK_KERNEL
-    } else if (paramb.charge_mode == 3) {
+    } 
+    
+    // mode 3 has real space only
+    if (paramb.charge_mode == 3) {
       find_force_charge_real_space_only<<<grid_size, block_size>>>(
         dataset[device_id].N,
         charge_para.alpha,

src/main_nep/parameters.cu

@@ -243,7 +243,7 @@ void Parameters::calculate_parameters()
     }
     std::vector<std::string> tokens;
     const int NUM89 = 89;
-    const int num_ann_per_element = (dim + 2) * num_neurons1;
+    const int num_ann_per_element = (dim + (charge_mode ? 3 : 2)) * num_neurons1;
     const int num_ann = NUM89 * num_ann_per_element + 1;
     const int num_cnk_radial = NUM89 * NUM89 * (n_max_radial + 1) * (basis_size_radial + 1);
     const int num_cnk_angular = NUM89 * NUM89 * (n_max_angular + 1) * (basis_size_angular + 1);

src/main_nep/snes.cu

@@ -130,7 +130,7 @@ void SNES::initialize_mu_and_sigma_fine_tune(Parameters& para)
   };
   // read in the whole foundation file first
   const int NUM89 = 89;
-  const int num_ann_per_element = (para.dim + 2) * para.num_neurons1;
+  const int num_ann_per_element = (para.dim + (para.charge_mode ? 3 : 2)) * para.num_neurons1;
   const int num_ann = NUM89 * num_ann_per_element + 1;
   const int num_cnk_radial = NUM89 * NUM89 * (para.n_max_radial + 1) * (para.basis_size_radial + 1);
   const int num_cnk_angular = NUM89 * NUM89 * (para.n_max_angular + 1) * (para.basis_size_angular + 1);

src/main_nep/structure.cu

@@ -215,6 +215,15 @@ static void read_one_structure(
     PRINT_INPUT_ERROR("'energy' is missing in the second line of a frame.");
   }
 
+  // get total charge (optional; default is 0)
+  for (const auto& token : tokens) {
+    const std::string charge_string = "charge=";
+    if (token.substr(0, charge_string.length()) == charge_string) {
+      structure.charge = get_double_from_token(
+        token.substr(charge_string.length(), token.length()), xyz_filename.c_str(), line_number);
+    }
+  }
+
   structure.has_temperature = false;
   for (const auto& token : tokens) {
     const std::string temperature_string = "temperature=";

tools/Miscellaneous/for_coding/for_perioidc_table/nep_data_toolkit.cpp

@@ -21,6 +21,7 @@
 #include <sstream>
 #include <string>
 #include <vector>
+#include <iomanip>
 
 static std::string remove_spaces_step1(const std::string& line)
 {
@@ -161,6 +162,7 @@ struct Structure {
   bool has_sid = false;
   bool has_virial = false;
   bool has_stress = false;
+  double charge = 0.0;
   double energy_weight = 1.0;
   double energy;
   double weight;
@@ -232,6 +234,15 @@ static void read_one_structure(std::ifstream& input, Structure& structure)
     }
   }
 
+  // get charge (optional)
+  for (const auto& token : tokens) {
+    const std::string charge_string = "charge=";
+    if (token.substr(0, charge_string.length()) == charge_string) {
+      structure.charge = get_double_from_token(
+        token.substr(charge_string.length(), token.length()), __FILE__, __LINE__);
+    }
+  }
+
   // get energy_weight (optional)
   for (const auto& token : tokens) {
     const std::string energy_weight_string = "energy_weight=";
@@ -408,6 +419,10 @@ static void write_one_structure(std::ofstream& output, const Structure& structur
 {
   output << structure.num_atom << "\n";
 
+  if (structure.charge != 0.0) {
+    output << "charge=" << structure.charge << " ";
+  }
+
   if (structure.energy_weight != 1.0) {
     output << "energy_weight=" << structure.energy_weight << " ";
   }
@@ -421,7 +436,7 @@ static void write_one_structure(std::ofstream& output, const Structure& structur
   }
   output << "\" ";
 
-  output << "energy=" << structure.energy << " ";
+  output << "energy=" << std::fixed << std::setprecision(6) << structure.energy << " ";
 
   if (structure.has_virial) {
     output << "virial=\"";