dump nep.restart when writing nep.txt

Author: elindgren

src/main_nep/fitness.cu

@@ -415,6 +415,19 @@ void Fitness::write_nep_txt(FILE* fid_nep, Parameters& para, float* elite)
   }
 }
 
+void Fitness::get_save_potential_label(Parameters& para, const int generation, std::string& label) {
+    if (para.save_potential_format == 1) {
+      time_t rawtime;
+      time(&rawtime);
+      struct tm* timeinfo = localtime(&rawtime);
+      char buffer[200];
+      strftime(buffer, sizeof(buffer), "nep_y%Y_m%m_d%d_h%H_m%M_s%S_generation", timeinfo);
+      label = std::string(buffer) + std::to_string(generation + 1);
+    } else {
+      label = "nep_gen" + std::to_string(generation + 1);
+    }
+}
+
 void Fitness::report_error(
   Parameters& para,
   const int generation,
@@ -462,16 +475,8 @@ void Fitness::report_error(
 
     if (0 == (generation + 1) % para.save_potential) {
       std::string filename;
-      if (para.save_potential_format == 1) {
-        time_t rawtime;
-        time(&rawtime);
-        struct tm* timeinfo = localtime(&rawtime);
-        char buffer[200];
-        strftime(buffer, sizeof(buffer), "nep_y%Y_m%m_d%d_h%H_m%M_s%S_generation", timeinfo);
-        filename = std::string(buffer) + std::to_string(generation + 1) + ".txt";
-      } else {
-        filename = "nep_gen" + std::to_string(generation + 1) + ".txt";
-      }
+      get_save_potential_label(para, generation, filename);
+      filename += ".txt";
 
       FILE* fid_nep = my_fopen(filename.c_str(), "w");
       write_nep_txt(fid_nep, para, elite);

src/main_nep/fitness.cuh

@@ -37,6 +37,7 @@ public:
     const float loss_L2,
     float* elite);
   void predict(Parameters& para, float* elite);
+  void get_save_potential_label(Parameters& para, const int generation, std::string& filename);
 
 protected:
   bool has_test_set = false;

src/main_nep/snes.cu

@@ -356,7 +356,15 @@ void SNES::compute(Parameters& para, Fitness* fitness_function)
 
       update_mu_and_sigma(para);
       if (0 == (n + 1) % 100) {
-        output_mu_and_sigma(para);
+        const char* filename = "nep.restart";
+        output_mu_and_sigma(para, filename);
+      }
+      // Optionally save the nep.restart file at the same time as save_potential
+      if (0 == (n + 1) % para.save_potential) {
+        std::string restart_file;
+        fitness_function->get_save_potential_label(para, n, restart_file);
+        restart_file += ".restart";
+        output_mu_and_sigma(para, restart_file.c_str());
       }
     }
   } else {
@@ -642,12 +650,12 @@ void SNES::update_mu_and_sigma(Parameters& para)
   GPU_CHECK_KERNEL;
 }
 
-void SNES::output_mu_and_sigma(Parameters& para)
+void SNES::output_mu_and_sigma(Parameters& para, const char* filename)
 {
   gpuSetDevice(0); // normally use GPU-0
   gpu_mu.copy_to_host(mu.data());
   gpu_sigma.copy_to_host(sigma.data());
-  FILE* fid_restart = my_fopen("nep.restart", "w");
+  FILE* fid_restart = my_fopen(filename, "w");
   for (int n = 0; n < number_of_variables; ++n) {
     fprintf(fid_restart, "%15.7e %15.7e\n", mu[n], sigma[n]);
   }

src/main_nep/snes.cuh

@@ -70,5 +70,5 @@ protected:
   void regularize_NEP4(Parameters& para);
   void sort_population(Parameters& para);
   void update_mu_and_sigma(Parameters& para);
-  void output_mu_and_sigma(Parameters& para);
+  void output_mu_and_sigma(Parameters& para, const char* filename);
 };