Merge pull request brucefan1983#1293 from brucefan1983/try_population_mix

lambda_q and lambda_ z

Author: brucefan1983

doc/nep/input_parameters/index.rst

@@ -28,6 +28,8 @@ Below you can find a listing of keywords for the ``nep.in`` input file.
    lambda_e
    lambda_f
    lambda_v
+   lambda_q
+   lambda_z
    atomic_v
    lambda_shear
    force_delta

doc/nep/input_parameters/lambda_q.rst

@@ -0,0 +1,15 @@
+.. _kw_lambda_q:
+.. index::
+   single: lambda_q (keyword in nep.in)
+
+:attr:`lambda_q`
+================
+
+This keyword sets the weight :math:`\lambda_q` of the loss term associated with the **total charge** (the charge for a whole structure) in the :ref:`loss function <nep_loss_function>`.
+The syntax is::
+
+  lambda_q <weight>
+
+Here, :attr:`<weight>` represents :math:`\lambda_q`, which must satisfy :math:`\lambda_q \geq 0` and defaults to :math:`\lambda_q = 0.5`.
+
+This keyword is only relevant for the qNEP models.

doc/nep/input_parameters/lambda_z.rst

@@ -0,0 +1,15 @@
+.. _kw_lambda_z:
+.. index::
+   single: lambda_z (keyword in nep.in)
+
+:attr:`lambda_z`
+================
+
+This keyword sets the weight :math:`\lambda_z` of the loss term associated with the **Born effective charge** in the :ref:`loss function <nep_loss_function>`.
+The syntax is::
+
+  lambda_z <weight>
+
+Here, :attr:`<weight>` represents :math:`\lambda_z`, which must satisfy :math:`\lambda_z \geq 0` and defaults to :math:`\lambda_z = 0.5`.
+
+This keyword is only relevant for the qNEP models.

src/main_nep/dataset.cu

@@ -976,32 +976,52 @@ std::vector<float> Dataset::get_rmse_charge(Parameters& para, int device_id)
     error_gpu.data());
   CHECK(gpuMemcpy(error_cpu.data(), error_gpu.data(), mem, gpuMemcpyDeviceToHost));
   for (int n = 0; n < Nc; ++n) {
-      float rmse_temp = error_cpu[n];
-      for (int t = 0; t < para.num_types + 1; ++t) {
-        if (has_type[t * Nc + n]) {
-          rmse_array[t] += rmse_temp;
-          count_array[t] += 1;
-        }
+    float rmse_temp = error_cpu[n];
+    for (int t = 0; t < para.num_types + 1; ++t) {
+      if (has_type[t * Nc + n]) {
+        rmse_array[t] += rmse_temp;
+        count_array[t] += 1;
       }
+    }
   }
 
-  if (para.has_bec) {
-    gpu_sum_bec_error<<<Nc, block_size, sizeof(float) * block_size>>>(
-      N,
-      Na.data(),
-      Na_sum.data(),
-      bec.data(),
-      bec_ref_gpu.data(),
-      error_gpu.data());
-    CHECK(gpuMemcpy(error_cpu.data(), error_gpu.data(), mem, gpuMemcpyDeviceToHost));
-    for (int n = 0; n < Nc; ++n) {
-      if (structures[n].has_bec) {
-        float rmse_temp = error_cpu[n];
-        for (int t = 0; t < para.num_types + 1; ++t) {
-          if (has_type[t * Nc + n]) {
-            rmse_array[t] += rmse_temp / (Na_cpu[n]);
-            count_array[t] += 9;
-          }
+  for (int t = 0; t <= para.num_types; ++t) {
+    if (count_array[t] > 0) {
+      rmse_array[t] = sqrt(rmse_array[t] / count_array[t]);
+    }
+  }
+  return rmse_array;
+}
+
+std::vector<float> Dataset::get_rmse_bec(Parameters& para, int device_id)
+{
+  std::vector<float> rmse_array(para.num_types + 1, 0.0f);
+  if (!(para.charge_mode && para.has_bec)) {
+    return rmse_array;
+  }
+
+  CHECK(gpuSetDevice(device_id));
+
+  std::vector<int> count_array(para.num_types + 1, 0);
+
+  int mem = sizeof(float) * Nc;
+  const int block_size = 256;
+
+  gpu_sum_bec_error<<<Nc, block_size, sizeof(float) * block_size>>>(
+    N,
+    Na.data(),
+    Na_sum.data(),
+    bec.data(),
+    bec_ref_gpu.data(),
+    error_gpu.data());
+  CHECK(gpuMemcpy(error_cpu.data(), error_gpu.data(), mem, gpuMemcpyDeviceToHost));
+  for (int n = 0; n < Nc; ++n) {
+    if (structures[n].has_bec) {
+      float rmse_temp = error_cpu[n];
+      for (int t = 0; t < para.num_types + 1; ++t) {
+        if (has_type[t * Nc + n]) {
+          rmse_array[t] += rmse_temp / (Na_cpu[n]);
+          count_array[t] += 9;
         }
       }
     }
@@ -1014,4 +1034,3 @@ std::vector<float> Dataset::get_rmse_charge(Parameters& para, int device_id)
   }
   return rmse_array;
 }
-

src/main_nep/dataset.cuh

@@ -92,6 +92,7 @@ public:
   std::vector<float> get_rmse_virial(Parameters& para, const bool use_weight, int device_id);
   std::vector<float> get_rmse_avirial(Parameters& para, const bool use_weight, int device_id);
   std::vector<float> get_rmse_charge(Parameters& para, int device_id);
+  std::vector<float> get_rmse_bec(Parameters& para, int device_id);
 
 private:
   void copy_structures(std::vector<Structure>& structures_input, int n1, int n2);

src/main_nep/fitness.cu

@@ -149,7 +149,14 @@ Fitness::~Fitness()
 }
 
 void Fitness::compute(
-  const int generation, Parameters& para, const float* population, float* fitness)
+  const int generation, 
+  Parameters& para, 
+  const float* population, 
+  float* fitness_energy,
+  float* fitness_force,
+  float* fitness_virial,
+  float* fitness_charge,
+  float* fitness_bec)
 {
   int deviceCount;
   CHECK(gpuGetDeviceCount(&deviceCount));
@@ -180,20 +187,19 @@ void Fitness::compute(
         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);
+        auto rmse_bec_array = train_set[batch_id][m].get_rmse_bec(para, m);
 
         for (int t = 0; t <= para.num_types; ++t) {
-          fitness[deviceCount * n + m + (7 * t + 3) * para.population_size] =
+          fitness_energy[deviceCount * n + m + t * para.population_size] =
             para.lambda_e * rmse_energy_array[t];
-          fitness[deviceCount * n + m + (7 * t + 4) * para.population_size] =
+          fitness_force[deviceCount * n + m + t * para.population_size] =
             para.lambda_f * rmse_force_array[t];
-          fitness[deviceCount * n + m + (7 * t + 5) * para.population_size] =
+          fitness_virial[deviceCount * n + m + t * para.population_size] =
             para.lambda_v * rmse_virial_array[t];
-          if (para.charge_mode) {
-            fitness[deviceCount * n + m + (7 * t + 6) * para.population_size] =
-              para.lambda_q * rmse_charge_array[t];
-          } else {
-            fitness[deviceCount * n + m + (7 * t + 6) * para.population_size] = 0.0f;
-          }
+          fitness_charge[deviceCount * n + m + t * para.population_size] =
+            para.lambda_q * rmse_charge_array[t];
+          fitness_bec[deviceCount * n + m + t * para.population_size] =
+            para.lambda_z * rmse_bec_array[t];
         }
       }
     }
@@ -216,33 +222,38 @@ void Fitness::compute(
             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);
+            auto rmse_bec_array = train_set[batch_id][m].get_rmse_bec(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];
+              float old_value = fitness_energy[deviceCount * n + m + t * para.population_size];
               float new_value = para.lambda_e * rmse_energy_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 + 3) * para.population_size] = new_value;
+              fitness_energy[deviceCount * n + m + t * para.population_size] = new_value;
               // force
-              old_value = fitness[deviceCount * n + m + (7 * t + 4) * para.population_size];
+              old_value = fitness_force[deviceCount * n + m + t * para.population_size];
               new_value = para.lambda_f * rmse_force_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 + 4) * para.population_size] = new_value;
+              fitness_force[deviceCount * n + m + t * para.population_size] = new_value;
               // virial
-              old_value = fitness[deviceCount * n + m + (7 * t + 5) * para.population_size];
+              old_value = fitness_virial[deviceCount * n + m + t * para.population_size];
               new_value = para.lambda_v * rmse_virial_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 + 5) * para.population_size] = new_value;
+              fitness_virial[deviceCount * n + m + t * 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;
-              }
+              old_value = fitness_charge[deviceCount * n + m + t * 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_charge[deviceCount * n + m + t * para.population_size] = new_value;
+              // BEC
+              old_value = fitness_bec[deviceCount * n + m + t * para.population_size];
+              new_value = para.lambda_z * rmse_bec_array[t];
+              new_value = old_value * old_value * count_batch + new_value * new_value;
+              new_value = sqrt(new_value / (count_batch + 1));
+              fitness_bec[deviceCount * n + m + t * para.population_size] = new_value;
             }
           }
         }
@@ -435,10 +446,14 @@ void Fitness::report_error(
       train_set[batch_id][0].get_rmse_energy(para, energy_shift_per_structure, false, true, 0);
     auto rmse_force_train_array = train_set[batch_id][0].get_rmse_force(para, false, 0);
     auto rmse_virial_train_array = train_set[batch_id][0].get_rmse_virial(para, false, 0);
+    auto rmse_charge_train_array = train_set[batch_id][0].get_rmse_charge(para, 0);
+    auto rmse_bec_train_array = train_set[batch_id][0].get_rmse_bec(para, 0);
 
     float rmse_energy_train = rmse_energy_train_array.back();
     float rmse_force_train = rmse_force_train_array.back();
     float rmse_virial_train = rmse_virial_train_array.back();
+    float rmse_charge_train = rmse_charge_train_array.back();
+    float rmse_bec_train = rmse_bec_train_array.back();
 
     // correct the last bias parameter in the NN
     if (para.train_mode == 0 || para.train_mode == 3) {
@@ -448,16 +463,22 @@ void Fitness::report_error(
     float rmse_energy_test = 0.0f;
     float rmse_force_test = 0.0f;
     float rmse_virial_test = 0.0f;
+    float rmse_charge_test = 0.0f;
+    float rmse_bec_test = 0.0f;
     if (has_test_set) {
       potential->find_force(para, elite, test_set, false, true, 1);
       float energy_shift_per_structure_not_used;
       auto rmse_energy_test_array =
         test_set[0].get_rmse_energy(para, energy_shift_per_structure_not_used, false, false, 0);
       auto rmse_force_test_array = test_set[0].get_rmse_force(para, false, 0);
       auto rmse_virial_test_array = test_set[0].get_rmse_virial(para, false, 0);
+      auto rmse_charge_test_array = test_set[0].get_rmse_charge(para, 0);
+      auto rmse_bec_test_array = test_set[0].get_rmse_bec(para, 0);
       rmse_energy_test = rmse_energy_test_array.back();
       rmse_force_test = rmse_force_test_array.back();
       rmse_virial_test = rmse_virial_test_array.back();
+      rmse_charge_test = rmse_charge_test_array.back();
+      rmse_bec_test = rmse_bec_test_array.back();
     }
 
     FILE* fid_nep = my_fopen("nep.txt", "w");
@@ -475,32 +496,71 @@ void Fitness::report_error(
     }
 
     if (para.train_mode == 0 || para.train_mode == 3) {
-      printf(
-        "%-8d%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f\n",
-        generation + 1,
-        loss_total,
-        loss_L1,
-        loss_L2,
-        rmse_energy_train,
-        rmse_force_train,
-        rmse_virial_train,
-        rmse_energy_test,
-        rmse_force_test,
-        rmse_virial_test);
-      fprintf(
-        fid_loss_out,
-        "%-8d%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f\n",
-        generation + 1,
-        loss_total,
-        loss_L1,
-        loss_L2,
-        rmse_energy_train,
-        rmse_force_train,
-        rmse_virial_train,
-        rmse_energy_test,
-        rmse_force_test,
-        rmse_virial_test);
+      if (!para.charge_mode) {
+        // NEP models
+        printf(
+          "%-8d%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f\n",
+          generation + 1,
+          loss_total,
+          loss_L1,
+          loss_L2,
+          rmse_energy_train,
+          rmse_force_train,
+          rmse_virial_train,
+          rmse_energy_test,
+          rmse_force_test,
+          rmse_virial_test);
+        fprintf(
+          fid_loss_out,
+          "%-8d%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f%-13.5f\n",
+          generation + 1,
+          loss_total,
+          loss_L1,
+          loss_L2,
+          rmse_energy_train,
+          rmse_force_train,
+          rmse_virial_train,
+          rmse_energy_test,
+          rmse_force_test,
+          rmse_virial_test);
+      } else {
+        // qNEP models:
+        printf(
+          "%-8d%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f\n",
+          generation + 1,
+          loss_total,
+          loss_L1,
+          loss_L2,
+          rmse_energy_train,
+          rmse_force_train,
+          rmse_virial_train,
+          rmse_charge_train,
+          rmse_bec_train,
+          rmse_energy_test,
+          rmse_force_test,
+          rmse_virial_test,
+          rmse_charge_test,
+          rmse_bec_test);
+        fprintf(
+          fid_loss_out,
+          "%-8d%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f%-9.5f\n",
+          generation + 1,
+          loss_total,
+          loss_L1,
+          loss_L2,
+          rmse_energy_train,
+          rmse_force_train,
+          rmse_virial_train,
+          rmse_charge_train,
+          rmse_bec_train,
+          rmse_energy_test,
+          rmse_force_test,
+          rmse_virial_test,
+          rmse_charge_test,
+          rmse_bec_test);
+      }
     } else {
+      // TNEP models:
       printf(
         "%-8d%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f\n",
         generation + 1,

src/main_nep/fitness.cuh

@@ -28,7 +28,7 @@ class Fitness
 public:
   Fitness(Parameters& para);
   ~Fitness();
-  void compute(const int generation, Parameters& para, const float*, float*);
+  void compute(const int generation, Parameters& para, const float*, float*, float*, float*, float*, float*);
   void report_error(
     Parameters& para,
     const int generation,