Merge pull request brucefan1983#907 from Yi-FanLi/avirial

Support fitting atomic virial, dipole, and polarizability

Author: brucefan1983

doc/nep/input_files/nep_in.rst

@@ -66,6 +66,8 @@ Keywords
      - weight of force loss term
    * - :ref:`lambda_v <kw_lambda_v>`
      - weight of virial loss term
+   * - :ref:`atomic_v <kw_atomic_v>`
+     - fit atomic or global virial
    * - :ref:`force_delta <kw_force_delta>`
      - bias term that can be used to make smaller forces more accurate
    * - :ref:`batch <kw_batch>`

doc/nep/input_parameters/atomic_v.rst

@@ -0,0 +1,20 @@
+.. _kw_atomic_v:
+.. index::
+   single: atomic_v (keyword in nep.in)
+
+:attr:`atomic_v`
+================
+
+This keyword sets the mode :math:`\atomic_v` of whether to fit atomic or global quantities for dipole (`model_type = 1`) or polarizability (`model_type = 2`). Only one of atomic and global can be fitted at a time. Fitting both simultaneously is not supported. For the virial tensor (`model_type = 0`), only the global model is supported. 
+The syntax is::
+
+  atomic_v <mode>
+
+where :attr:`<mode>` must be an integer that can assume one of the following values.
+
+=====  ===========================
+Value  Mode 
+-----  ---------------------------
+0      fit global tensor (default)
+1      fit atomic tensor
+=====  ===========================

src/main_nep/dataset.cu

@@ -21,6 +21,8 @@
 #include "utilities/gpu_macro.cuh"
 #include "utilities/nep_utilities.cuh"
 #include <cstring>
+#include <iostream>
+#include <stdexcept>
 
 void Dataset::copy_structures(std::vector<Structure>& structures_input, int n1, int n2)
 {
@@ -32,6 +34,8 @@ void Dataset::copy_structures(std::vector<Structure>& structures_input, int n1,
     structures[n].num_atom = structures_input[n_input].num_atom;
     structures[n].weight = structures_input[n_input].weight;
     structures[n].has_virial = structures_input[n_input].has_virial;
+    structures[n].has_atomic_virial = structures_input[n_input].has_atomic_virial;
+    structures[n].atomic_virial_diag_only = structures_input[n_input].atomic_virial_diag_only;
     structures[n].charge = structures_input[n_input].charge;
     structures[n].energy = structures_input[n_input].energy;
     structures[n].energy_weight = structures_input[n_input].energy_weight;
@@ -68,6 +72,33 @@ void Dataset::copy_structures(std::vector<Structure>& structures_input, int n1,
       structures[n].fy[na] = structures_input[n_input].fy[na];
       structures[n].fz[na] = structures_input[n_input].fz[na];
     }
+
+    if (structures[n].has_atomic_virial != structures[0].has_atomic_virial) {
+      throw std::runtime_error("All structures must have the same has_atomic_virial flag.");
+    }
+    if (structures[n].atomic_virial_diag_only != structures[0].atomic_virial_diag_only) {
+      throw std::runtime_error("All structures must have the same atomic_virial_diag_only flag.");
+    }
+    if (structures[n].has_atomic_virial) {
+      structures[n].avirialxx.resize(structures[n].num_atom);
+      structures[n].avirialyy.resize(structures[n].num_atom);
+      structures[n].avirialzz.resize(structures[n].num_atom);
+      for (int na = 0; na < structures[n].num_atom; ++na) {
+        structures[n].avirialxx[na] = structures_input[n_input].avirialxx[na];
+        structures[n].avirialyy[na] = structures_input[n_input].avirialyy[na];
+        structures[n].avirialzz[na] = structures_input[n_input].avirialzz[na];
+      }
+      if (!structures[n].atomic_virial_diag_only) {
+        structures[n].avirialxy.resize(structures[n].num_atom);
+        structures[n].avirialyz.resize(structures[n].num_atom);
+        structures[n].avirialzx.resize(structures[n].num_atom);
+        for (int na = 0; na < structures[n].num_atom; ++na) {
+          structures[n].avirialxy[na] = structures_input[n_input].avirialxy[na];
+          structures[n].avirialyz[na] = structures_input[n_input].avirialyz[na];
+          structures[n].avirialzx[na] = structures_input[n_input].avirialzx[na];
+        }
+      }
+    }
   }
 }
 
@@ -142,6 +173,9 @@ void Dataset::initialize_gpu_data(Parameters& para)
   energy_weight_cpu.resize(Nc);
   virial_ref_cpu.resize(Nc * 6);
   force_ref_cpu.resize(N * 3);
+  if (structures[0].has_atomic_virial) {
+    avirial_ref_cpu.resize(N * (structures[0].atomic_virial_diag_only ? 3 : 6));
+  }
   temperature_ref_cpu.resize(N);
 
   for (int n = 0; n < Nc; ++n) {
@@ -170,6 +204,16 @@ void Dataset::initialize_gpu_data(Parameters& para)
       force_ref_cpu[Na_sum_cpu[n] + na + N] = structures[n].fy[na];
       force_ref_cpu[Na_sum_cpu[n] + na + N * 2] = structures[n].fz[na];
       temperature_ref_cpu[Na_sum_cpu[n] + na] = structures[n].temperature;
+      if (structures[n].has_atomic_virial) {
+        avirial_ref_cpu[Na_sum_cpu[n] + na] = structures[n].avirialxx[na];
+        avirial_ref_cpu[Na_sum_cpu[n] + na + N] = structures[n].avirialyy[na];
+        avirial_ref_cpu[Na_sum_cpu[n] + na + N * 2] = structures[n].avirialzz[na];
+        if (!structures[n].atomic_virial_diag_only) {
+          avirial_ref_cpu[Na_sum_cpu[n] + na + N * 3] = structures[n].avirialxy[na];
+          avirial_ref_cpu[Na_sum_cpu[n] + na + N * 4] = structures[n].avirialyz[na];
+          avirial_ref_cpu[Na_sum_cpu[n] + na + N * 5] = structures[n].avirialzx[na];
+        }
+      }
     }
   }
 
@@ -179,13 +223,19 @@ void Dataset::initialize_gpu_data(Parameters& para)
   energy_weight_gpu.resize(Nc);
   virial_ref_gpu.resize(Nc * 6);
   force_ref_gpu.resize(N * 3);
+  if (structures[0].has_atomic_virial) {
+    avirial_ref_gpu.resize(N * (structures[0].atomic_virial_diag_only ? 3 : 6));
+  }
   temperature_ref_gpu.resize(N);
   type_weight_gpu.copy_from_host(para.type_weight_cpu.data());
   charge_ref_gpu.copy_from_host(charge_ref_cpu.data());
   energy_ref_gpu.copy_from_host(energy_ref_cpu.data());
   energy_weight_gpu.copy_from_host(energy_weight_cpu.data());
   virial_ref_gpu.copy_from_host(virial_ref_cpu.data());
   force_ref_gpu.copy_from_host(force_ref_cpu.data());
+  if (structures[0].has_atomic_virial) {
+    avirial_ref_gpu.copy_from_host(avirial_ref_cpu.data());
+  }
   temperature_ref_gpu.copy_from_host(temperature_ref_cpu.data());
 
   box.resize(Nc * 18);
@@ -447,6 +497,157 @@ std::vector<float> Dataset::get_rmse_force(Parameters& para, const bool use_weig
   return rmse_array;
 }
 
+static __global__ void gpu_sum_avirial_diag_only_error(
+  const int N,
+  int* g_Na,
+  int* g_Na_sum,
+  int* g_type,
+  float* g_type_weight,
+  float* g_virial,
+  float* g_avxx_ref,
+  float* g_avyy_ref,
+  float* g_avzz_ref,
+  float* error_gpu)
+{
+  int tid = threadIdx.x;
+  int bid = blockIdx.x;
+  int N1 = g_Na_sum[bid];
+  int N2 = N1 + g_Na[bid];
+  extern __shared__ float s_error[];
+  s_error[tid] = 0.0f;
+
+  for (int n = N1 + tid; n < N2; n += blockDim.x) {
+    float avxx_ref = g_avxx_ref[n];
+    float avyy_ref = g_avyy_ref[n];
+    float avzz_ref = g_avzz_ref[n];
+    float dxx = g_virial[n] - avxx_ref;
+    float dyy = g_virial[1 * N + n] - avyy_ref;
+    float dzz = g_virial[2 * N + n] - avzz_ref;
+    float diff_square = dxx * dxx + dyy * dyy + dzz * dzz;
+    s_error[tid] += diff_square;
+  }
+  __syncthreads();
+
+  for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
+    if (tid < offset) {
+      s_error[tid] += s_error[tid + offset];
+    }
+    __syncthreads();
+  }
+
+  if (tid == 0) {
+    error_gpu[bid] = s_error[0];
+  }
+}
+
+static __global__ void gpu_sum_avirial_error(
+  const int N,
+  int* g_Na,
+  int* g_Na_sum,
+  int* g_type,
+  float* g_type_weight,
+  float* g_virial,
+  float* g_avxx_ref,
+  float* g_avyy_ref,
+  float* g_avzz_ref,
+  float* g_avxy_ref,
+  float* g_avyz_ref,
+  float* g_avzx_ref,
+  float* error_gpu)
+{
+  int tid = threadIdx.x;
+  int bid = blockIdx.x;
+  int N1 = g_Na_sum[bid];
+  int N2 = N1 + g_Na[bid];
+  extern __shared__ float s_error[];
+  s_error[tid] = 0.0f;
+
+  for (int n = N1 + tid; n < N2; n += blockDim.x) {
+    float avxx_ref = g_avxx_ref[n];
+    float avyy_ref = g_avyy_ref[n];
+    float avzz_ref = g_avzz_ref[n];
+    float avxy_ref = g_avxy_ref[n];
+    float avyz_ref = g_avyz_ref[n];
+    float avzx_ref = g_avzx_ref[n];
+    float dxx = g_virial[n] - avxx_ref;
+    float dyy = g_virial[1 * N + n] - avyy_ref;
+    float dzz = g_virial[2 * N + n] - avzz_ref;
+    float dxy = g_virial[3 * N + n] - avxy_ref;
+    float dyz = g_virial[4 * N + n] - avyz_ref;
+    float dzx = g_virial[5 * N + n] - avzx_ref;
+    float diff_square = dxx * dxx + dyy * dyy + dzz * dzz + dxy * dxy + dyz * dyz + dzx * dzx;
+    s_error[tid] += diff_square;
+  }
+  __syncthreads();
+
+  for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
+    if (tid < offset) {
+      s_error[tid] += s_error[tid + offset];
+    }
+    __syncthreads();
+  }
+
+  if (tid == 0) {
+    error_gpu[bid] = s_error[0];
+  }
+}
+
+std::vector<float> Dataset::get_rmse_avirial(Parameters& para, const bool use_weight, int device_id)
+{
+  CHECK(gpuSetDevice(device_id));
+  const int block_size = 256;
+
+  if (structures[0].atomic_virial_diag_only) {
+    gpu_sum_avirial_diag_only_error<<<Nc, block_size, sizeof(float) * block_size>>>(
+      N,
+      Na.data(),
+      Na_sum.data(),
+      type.data(),
+      type_weight_gpu.data(),
+      virial.data(),
+      avirial_ref_gpu.data(),
+      avirial_ref_gpu.data() + N,
+      avirial_ref_gpu.data() + N * 2,
+      error_gpu.data());
+  } else {
+    gpu_sum_avirial_error<<<Nc, block_size, sizeof(float) * block_size>>>(
+      N,
+      Na.data(),
+      Na_sum.data(),
+      type.data(),
+      type_weight_gpu.data(),
+      virial.data(),
+      avirial_ref_gpu.data(),
+      avirial_ref_gpu.data() + N,
+      avirial_ref_gpu.data() + N * 2,
+      avirial_ref_gpu.data() + N * 3,
+      avirial_ref_gpu.data() + N * 4,
+      avirial_ref_gpu.data() + N * 5,
+      error_gpu.data());
+  }
+  int mem = sizeof(float) * Nc;
+  CHECK(gpuMemcpy(error_cpu.data(), error_gpu.data(), mem, gpuMemcpyDeviceToHost));
+
+  std::vector<float> rmse_array(para.num_types + 1, 0.0f);
+  std::vector<int> count_array(para.num_types + 1, 0);
+  for (int n = 0; n < Nc; ++n) {
+    float rmse_temp = use_weight ? weight_cpu[n] * weight_cpu[n] * error_cpu[n] : 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] += Na_cpu[n];
+      }
+    }
+  }
+
+  for (int t = 0; t <= para.num_types; ++t) {
+    if (count_array[t] > 0) {
+      rmse_array[t] = sqrt(rmse_array[t] / (count_array[t] * 6));
+    }
+  }
+  return rmse_array;
+}
+
 static __global__ void
 gpu_get_energy_shift(
   int* g_Na, 
@@ -625,6 +826,9 @@ static __global__ void gpu_sum_virial_error(
 
 std::vector<float> Dataset::get_rmse_virial(Parameters& para, const bool use_weight, int device_id)
 {
+  if (para.atomic_v) {
+    return get_rmse_avirial(para, use_weight, device_id);
+  }
   CHECK(gpuSetDevice(device_id));
 
   std::vector<float> rmse_array(para.num_types + 1, 0.0f);

src/main_nep/dataset.cuh

@@ -43,22 +43,26 @@ public:
   GPU_Vector<float> energy;      // calculated energy in GPU
   GPU_Vector<float> virial;      // calculated virial in GPU
   GPU_Vector<float> force;       // calculated force in GPU
+  GPU_Vector<float> avirial; // calculated atomic virial in GPU
   std::vector<float> charge_cpu;  // calculated charge in CPU
   std::vector<float> energy_cpu; // calculated energy in CPU
   std::vector<float> virial_cpu; // calculated virial in CPU
   std::vector<float> force_cpu;  // calculated force in CPU
+  std::vector<float> avirial_cpu;   // calculated atomic virial in CPU
 
   GPU_Vector<float> energy_weight_gpu;    // energy weight in GPU
   GPU_Vector<float> charge_ref_gpu;       // reference charge in GPU
   GPU_Vector<float> energy_ref_gpu;       // reference energy in GPU
   GPU_Vector<float> virial_ref_gpu;       // reference virial in GPU
   GPU_Vector<float> force_ref_gpu;        // reference force in GPU
+  GPU_Vector<float> avirial_ref_gpu;     // reference atomic virial in GPU
   GPU_Vector<float> temperature_ref_gpu;  // reference temperature in GPU
   std::vector<float> energy_weight_cpu;   // energy weight in CPU
   std::vector<float> charge_ref_cpu;      // reference charge in CPU
   std::vector<float> energy_ref_cpu;      // reference energy in CPU
   std::vector<float> virial_ref_cpu;      // reference virial in CPU
   std::vector<float> force_ref_cpu;       // reference force in CPU
+  std::vector<float> avirial_ref_cpu;      // reference atomic virial in CPU
   std::vector<float> weight_cpu;          // configuration weight in CPU
   std::vector<float> temperature_ref_cpu; // reference temeprature in CPU
 
@@ -81,6 +85,7 @@ public:
     const bool do_shift,
     int device_id);
   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);
 
 private: