flip charge upon restarting

Author: brucefan1983

src/main_nep/parameters.cu

@@ -1299,15 +1299,23 @@ void Parameters::parse_charge_mode(const char** param, int num_param)
 {
   is_charge_mode_set = true;
 
-  if (num_param != 2) {
-    PRINT_INPUT_ERROR("charge_mode should have one parameter.\n");
+  if (num_param != 2 && num_param != 3) {
+    PRINT_INPUT_ERROR("charge_mode should have one or two parameters.\n");
   }
   if (!is_valid_int(param[1], &charge_mode)) {
     PRINT_INPUT_ERROR("charge mode should be an integer.\n");
   }
   if (charge_mode < 0 || charge_mode > 5) {
     PRINT_INPUT_ERROR("charge mode should be 0 or 1 or 2 or 3 or 4 or 5.");
   }
+  if (num_param == 3) {
+    if (!is_valid_int(param[2], &flip_charge)) {
+      PRINT_INPUT_ERROR("flip_charge should be an integer.\n");
+    }
+    if (flip_charge < 0 || flip_charge > 1) {
+      PRINT_INPUT_ERROR("flip_charge should be 0 or 1.");
+    }
+  }
 }
 
 void Parameters::parse_fine_tune(const char** param, int num_param)

src/main_nep/parameters.cuh

@@ -68,6 +68,7 @@ public:
   int output_descriptor;
   int charge_mode; // add dynamic charge to NEP potential model
   bool has_bec = false; // check if there are target BEC values
+  int flip_charge = 0; // 1 for flipping charges upon restarting
   int fine_tune = 0; // fine_tune option; 0=no, 1=yes
   std::string fine_tune_nep_txt = "";
   std::string fine_tune_nep_restart = "";

src/main_nep/snes.cu

@@ -121,6 +121,15 @@ void SNES::initialize_mu_and_sigma(Parameters& para)
       int count = fscanf(fid_restart, "%f%f", &mu[n], &sigma[n]);
       PRINT_SCANF_ERROR(count, 2, "Reading error for nep.restart.");
     }
+    // flip the charges if needed
+    if (para.charge_mode && para.flip_charge) {
+      const int num_part = (para.dim + 2) * para.num_neurons1;
+      for (int t = 0; t < para.num_types; ++t) {
+        for (int n = para.number_of_variables_ann_1 * t + num_part; n < para.number_of_variables_ann_1 * (t + 1); ++n) {
+          mu[n] = -mu[n];
+        }
+      }
+    }
     fclose(fid_restart);
   }
   gpuSetDevice(0); // normally use GPU-0