Merge pull request brucefan1983#1036 from brucefan1983/charge_loss

enable charge loss term

Author: brucefan1983

src/main_nep/dataset.cu

@@ -159,6 +159,7 @@ void Dataset::initialize_gpu_data(Parameters& para)
   std::vector<int> type_cpu(N);
 
   charge.resize(N);
+  charge_shifted.resize(N);
   energy.resize(N);
   virial.resize(N * 6);
   force.resize(N * 3);

src/main_nep/dataset.cuh

@@ -40,6 +40,7 @@ public:
   GPU_Vector<int> num_cell;       // number of cells in the expanded box (3 components)
 
   GPU_Vector<float> charge;      // calculated charge in GPU
+  GPU_Vector<float> charge_shifted;      // shifted charge in GPU
   GPU_Vector<float> energy;      // calculated energy in GPU
   GPU_Vector<float> virial;      // calculated virial in GPU
   GPU_Vector<float> force;       // calculated force in GPU

src/main_nep/fitness.cu

@@ -215,6 +215,7 @@ void Fitness::compute(
               para, energy_shift_per_structure_not_used, true, true, m);
             auto rmse_force_array = train_set[batch_id][m].get_rmse_force(para, true, m);
             auto rmse_virial_array = train_set[batch_id][m].get_rmse_virial(para, true, m);
+            auto rmse_charge_array = train_set[batch_id][m].get_rmse_charge(para, m);
             for (int t = 0; t <= para.num_types; ++t) {
               // energy
               float old_value = fitness[deviceCount * n + m + (7 * t + 3) * para.population_size];
@@ -234,6 +235,14 @@ void Fitness::compute(
               new_value = old_value * old_value * count_batch + new_value * new_value;
               new_value = sqrt(new_value / (count_batch + 1));
               fitness[deviceCount * n + m + (7 * t + 5) * para.population_size] = new_value;
+              // charge
+              if (para.charge_mode) {
+                old_value = fitness[deviceCount * n + m + (7 * t + 6) * para.population_size];
+                new_value = para.lambda_q * rmse_charge_array[t];
+                new_value = old_value * old_value * count_batch + new_value * new_value;
+                new_value = sqrt(new_value / (count_batch + 1));
+                fitness[deviceCount * n + m + (7 * t + 6) * para.population_size] = new_value;
+              }
             }
           }
         }

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)
+  float* g_charge,
+  float* g_charge_shifted)
 {
   int tid = threadIdx.x;
   int N1 = Na_sum[blockIdx.x];
@@ -1102,7 +1103,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[n] -= s_charge[0] / (N2 - N1);
+      g_charge_shifted[n] = g_charge[n] - s_charge[0] / (N2 - N1);
     }
   }
 }
@@ -1248,7 +1249,8 @@ void NEP_Charge::find_force(
     zero_total_charge<<<dataset[device_id].Nc, 1024>>>(
       dataset[device_id].Na.data(),
       dataset[device_id].Na_sum.data(),
-      dataset[device_id].charge.data());
+      dataset[device_id].charge.data(),
+      dataset[device_id].charge_shifted.data());
     GPU_CHECK_KERNEL
 
     if (paramb.charge_mode != 3) {
@@ -1269,7 +1271,7 @@ void NEP_Charge::find_force(
         charge_para.num_kpoints_max,
         dataset[device_id].Na.data(),
         dataset[device_id].Na_sum.data(),
-        dataset[device_id].charge.data(),
+        dataset[device_id].charge_shifted.data(),
         dataset[device_id].r.data(),
         dataset[device_id].r.data() + dataset[device_id].N,
         dataset[device_id].r.data() + dataset[device_id].N * 2,
@@ -1287,7 +1289,7 @@ void NEP_Charge::find_force(
         charge_para.alpha_factor,
         dataset[device_id].Na.data(),
         dataset[device_id].Na_sum.data(),
-        dataset[device_id].charge.data(),
+        dataset[device_id].charge_shifted.data(),
         dataset[device_id].r.data(),
         dataset[device_id].r.data() + dataset[device_id].N,
         dataset[device_id].r.data() + dataset[device_id].N * 2,
@@ -1316,7 +1318,7 @@ void NEP_Charge::find_force(
         charge_para.two_alpha_over_sqrt_pi,
         nep_data[device_id].NN_radial.data(),
         nep_data[device_id].NL_radial.data(),
-        dataset[device_id].charge.data(),
+        dataset[device_id].charge_shifted.data(),
         nep_data[device_id].x12_radial.data(),
         nep_data[device_id].y12_radial.data(),
         nep_data[device_id].z12_radial.data(),
@@ -1336,7 +1338,7 @@ void NEP_Charge::find_force(
         charge_para.B,
         nep_data[device_id].NN_radial.data(),
         nep_data[device_id].NL_radial.data(),
-        dataset[device_id].charge.data(),
+        dataset[device_id].charge_shifted.data(),
         nep_data[device_id].x12_radial.data(),
         nep_data[device_id].y12_radial.data(),
         nep_data[device_id].z12_radial.data(),