report charge and bec rmse

Author: brucefan1983

src/main_nep/fitness.cu

@@ -446,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) {
@@ -459,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");
@@ -486,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%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f%-13.5f%-13.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_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%-11.5f%-11.5f%-11.5f%-13.5f%-13.5f%-13.5f%-13.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_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/snes.cu

@@ -318,18 +318,37 @@ void SNES::compute(Parameters& para, Fitness* fitness_function)
   if (para.prediction == 0) {
 
     if (para.train_mode == 0 || para.train_mode == 3) {
-      printf(
-        "%-8s%-11s%-11s%-11s%-13s%-13s%-13s%-13s%-13s%-13s\n",
-        "Step",
-        "Total-Loss",
-        "L1Reg-Loss",
-        "L2Reg-Loss",
-        "RMSE-E-Train",
-        "RMSE-F-Train",
-        "RMSE-V-Train",
-        "RMSE-E-Test",
-        "RMSE-F-Test",
-        "RMSE-V-Test");
+      if (!para.charge_mode) {
+        printf(
+          "%-8s%-11s%-11s%-11s%-13s%-13s%-13s%-13s%-13s%-13s\n",
+          "Step",
+          "Total-Loss",
+          "L1Reg-Loss",
+          "L2Reg-Loss",
+          "RMSE-E-Train",
+          "RMSE-F-Train",
+          "RMSE-V-Train",
+          "RMSE-E-Test",
+          "RMSE-F-Test",
+          "RMSE-V-Test");
+      } else {
+        printf(
+          "%-8s%-11s%-11s%-11s%-13s%-13s%-13s%-13s%-13s%-13s%-13s%-13s%-13s%-13s\n",
+          "Step",
+          "Total-Loss",
+          "L1Reg-Loss",
+          "L2Reg-Loss",
+          "RMSE-E-Train",
+          "RMSE-F-Train",
+          "RMSE-V-Train",
+          "RMSE-Q-Train",
+          "RMSE-Z-Train",
+          "RMSE-E-Test",
+          "RMSE-F-Test",
+          "RMSE-V-Test",
+          "RMSE-Q-Test",
+          "RMSE-Z-Test");
+      }
     } else {
       printf(
         "%-8s%-11s%-11s%-11s%-13s%-13s\n",