Merge pull request brucefan1983#1169 from brucefan1983/improve_charge_train

Improve charge train

Author: brucefan1983

src/main_nep/parameters.cu

@@ -214,12 +214,16 @@ void Parameters::calculate_parameters()
   }
 
   if (version == 3) {
+    number_of_variables_ann_1 = (dim + 2) * num_neurons1;
     number_of_variables_ann = (dim + 2) * num_neurons1 + 1;
   } else if (version == 4) {
+    number_of_variables_ann_1 = (dim + 2) * num_neurons1;
     number_of_variables_ann = (dim + 2) * num_neurons1 * num_types + 1;
     if (charge_mode) {
+      number_of_variables_ann_1 += num_neurons1;
       number_of_variables_ann += num_neurons1 * num_types + 1;
       if (charge_mode == 4) {
+        number_of_variables_ann_1 += num_neurons1;
         number_of_variables_ann += num_neurons1 * num_types;
       }
     }
@@ -258,14 +262,7 @@ void Parameters::calculate_parameters()
     }
     std::vector<std::string> tokens;
     const int NUM89 = 89;
-    int num_outputs = 1;
-    if (charge_mode >= 1 && charge_mode <= 3) {
-      num_outputs = 2;
-    } else if (charge_mode == 4) {
-      num_outputs = 3;
-    }
-    const int num_ann_per_element = (dim + 1 + num_outputs) * num_neurons1;
-    const int num_ann = NUM89 * num_ann_per_element + (charge_mode ? 2 : 1);
+    const int num_ann = NUM89 * number_of_variables_ann_1 + (charge_mode ? 2 : 1);
     const int num_cnk_radial = NUM89 * NUM89 * (n_max_radial + 1) * (basis_size_radial + 1);
     const int num_cnk_angular = NUM89 * NUM89 * (n_max_angular + 1) * (basis_size_angular + 1);
     const int num_tot = num_ann + num_cnk_radial + num_cnk_angular;

src/main_nep/parameters.cuh

@@ -106,6 +106,7 @@ public:
   int dim_angular;                    // number of angular descriptor components
   int number_of_variables;            // total number of parameters (NN and descriptor)
   int number_of_variables_ann;        // number of parameters in the ANN only
+  int number_of_variables_ann_1;      // number of parameters in the ANN for one element
   int number_of_variables_descriptor; // number of parameters in the descriptor only
 
   // some arrays

src/main_nep/snes.cu

@@ -108,20 +108,13 @@ void SNES::initialize_mu_and_sigma(Parameters& para)
     }
     // make sure the initial charges are zero
     if (para.charge_mode) {
-      int num_outputs = 1;
-      if (para.charge_mode >= 1 && para.charge_mode <= 3) {
-        num_outputs = 2;
-      } else if (para.charge_mode == 4) {
-        num_outputs = 3;
-      }
-      const int num_full = (para.dim + 1 + num_outputs) * para.num_neurons1;
       const int num_part = (para.dim + 2) * para.num_neurons1;
       for (int t = 0; t < para.num_types; ++t) {
-        for (int n = num_full * t + num_part; n < num_full * (t + 1); ++n) {
+        for (int n = para.number_of_variables_ann_1 * t + num_part; n < para.number_of_variables_ann_1 * (t + 1); ++n) {
           mu[n] = 0.0f;
         }
       }
-      mu[num_full * para.num_types] = 2.0f; // make sure initial sqrt(epsilon_inf) > 0
+      mu[para.number_of_variables_ann_1 * para.num_types] = 2.0f; // make sure initial sqrt(epsilon_inf) > 0
     }
   } else {
     for (int n = 0; n < number_of_variables; ++n) {
@@ -147,14 +140,7 @@ void SNES::initialize_mu_and_sigma_fine_tune(Parameters& para)
   };
   // read in the whole foundation file first
   const int NUM89 = 89;
-  int num_outputs = 1;
-  if (para.charge_mode >= 1 && para.charge_mode <= 3) {
-    num_outputs = 2;
-  } else if (para.charge_mode == 4) {
-    num_outputs = 3;
-  }
-  const int num_ann_per_element = (para.dim + 1 + num_outputs) * para.num_neurons1;
-  const int num_ann = NUM89 * num_ann_per_element + (para.charge_mode ? 2 : 1);
+  const int num_ann = NUM89 * para.number_of_variables_ann_1 + (para.charge_mode ? 2 : 1);
   const int num_cnk_radial = NUM89 * NUM89 * (para.n_max_radial + 1) * (para.basis_size_radial + 1);
   const int num_cnk_angular = NUM89 * NUM89 * (para.n_max_angular + 1) * (para.basis_size_angular + 1);
   const int num_tot = num_ann + num_cnk_radial + num_cnk_angular;
@@ -182,9 +168,9 @@ void SNES::initialize_mu_and_sigma_fine_tune(Parameters& para)
   int count = 0;
   for (int i = 0; i < para.num_types; ++ i) {
     int element_index = element_map[para.atomic_numbers[i] - 1];
-    for (int j = 0; j < num_ann_per_element; ++j) {
-      mu[count] = restart_mu[element_index * num_ann_per_element + j];
-      sigma[count] = restart_sigma[element_index * num_ann_per_element + j];
+    for (int j = 0; j < para.number_of_variables_ann_1; ++j) {
+      mu[count] = restart_mu[element_index * para.number_of_variables_ann_1 + j];
+      sigma[count] = restart_sigma[element_index * para.number_of_variables_ann_1 + j];
       ++count;
     }
   }
@@ -252,30 +238,22 @@ void SNES::calculate_utility()
 
 void SNES::find_type_of_variable(Parameters& para)
 {
-  int num_outputs = 1;
-  if (para.charge_mode >= 1 && para.charge_mode <= 3) {
-    num_outputs = 2;
-  } else if (para.charge_mode == 4) {
-     num_outputs = 3;
-  }
-  int num_para_ann_per_type = (para.dim + 1 + num_outputs) * para.num_neurons1;
-
   int offset = 0;
 
   // NN part
   if (para.version != 3) {
     int num_ann = (para.train_mode == 2) ? 2 : 1;
     for (int ann = 0; ann < num_ann; ++ann) {
       for (int t = 0; t < para.num_types; ++t) {
-        for (int n = 0; n < num_para_ann_per_type; ++n) {
+        for (int n = 0; n < para.number_of_variables_ann_1; ++n) {
           type_of_variable[n + offset] = t;
         }
-        offset += num_para_ann_per_type;
+        offset += para.number_of_variables_ann_1;
       }
       offset += para.charge_mode ? 2 : 1; // the bias
     }
   } else {
-    offset += num_para_ann_per_type + 1;
+    offset += para.number_of_variables_ann_1 + 1;
   }
 
   // descriptor part
@@ -473,7 +451,8 @@ static __global__ void gpu_find_L1_L2_NEP4(
   s_cost_L2reg[tid] = 0.0f;
   for (int v = tid; v < number_of_variables; v += blockDim.x) {
     const float para = g_population[bid * number_of_variables + v];
-    if (g_type_of_variable[v] == g_type || g_type == g_num_types) {
+    if ((g_type_of_variable[v] == g_type) && (g_type != g_num_types) || 
+        (g_type_of_variable[v] != g_type) && (g_type == g_num_types))  {
       s_cost_L1reg[tid] += abs(para);
       s_cost_L2reg[tid] += para * para;
     }

src/makefile

@@ -83,7 +83,7 @@ HEADERS =                         \
 ###########################################################
 # executables
 ###########################################################
-all: gpumd nep gnep
+all: gpumd nep 
 gpumd: $(OBJ_GPUMD)
 	$(CC) $(LDFLAGS) $^ -o $@ $(LIBS)
 	@echo =================================================