Support fitting atomic virial, dipole, and polarizability #907

Yi-FanLi · 2025-02-14T15:37:12Z

Summary

This PR supports fitting atomic virial (for potential model), atomic dipole (for dipole model), and atomic polarizability (for polarizability model). The atmoic tensorial quatities can be used either alone or together with the global counterparts.

Modification

In struct Structure, 6 member variables (which are vectors of float) avirialxx, avirialyy, avirialzz, avirialxy, avirialyz, and avirialzx are added. These vectors are used to store the values of atomic virial and polarizability (use all 6 vectors) or dipole (only uses the first 3 vectors).

In class Dataset, a member function get_rmse_avirial is added. This function calculates the rmse of the atomic tensors. Two cuda kernels are added, i.e., gpu_sum_avirial_error and gpu_sum_avirial_diag_only_error. They are used to deal with the full 6-element tensor and the 3-element tensor with only diagonal terms, respectively.

In class Parameters, a member variable atomic_v is added. This variable can be set in nep.in and is used to determine whether to fit atomic or global tensor.

In class Fitness, the functions update_energy_force_virial, update_dipole, and update_polarizability are passed with a new variable para.atomic_v, which is used to decide whether to print atomic or global tensor.

Others

An example of training the atomic dipole of water is added in examples/15_NEP_atomic_dipole_water. The data is taken from the example of DeePMD-kit. The result of the training is as follows:

An example of training the atomic polarizability of water is added in examples/16_NEP_atomic_polarizability_water. The data is taken from the example of DeePMD-kit. The result of the training is as follows:

…ility

… been stored in virial

…update_polarizability

elindgren

Nice! 🚀 I had a few questions and comments. Additionally, I have some general questions:

When does one want to train on per-atom virials/dipoles/polarizabilities? Do they make sense in their per-atom versions?
Does the models take longer to train when using atomic vs global virials?

elindgren · 2025-02-21T15:33:19Z

src/main_nep/dataset.cu

+  CHECK(gpuSetDevice(device_id));
+  const int block_size = 256;
+
+  if (structures[0].atomic_virial_diag_only) {


Is there a check when loading the data to make sure that all structures have atomic_virial_diag_only? Or is there a change that these may be set differently for different structures?

elindgren · 2025-02-21T15:37:43Z

src/main_nep/fitness.cu

+    for (int nc = 0; nc < dataset.Nc; ++nc) {
+      int offset = dataset.Na_sum_cpu[nc];
+      for (int m = 0; m < dataset.structures[nc].num_atom; ++m) {
+        int n = offset + m;
+        fprintf(
+          fid_virial,
+          "%g %g %g %g %g %g %g %g %g %g %g %g\n",
+          dataset.virial_cpu[n] / dataset.Na_cpu[nc],
+          dataset.virial_cpu[n + dataset.N] / dataset.Na_cpu[nc],
+          dataset.virial_cpu[n + dataset.N * 2] / dataset.Na_cpu[nc],
+          dataset.virial_cpu[n + dataset.N * 3] / dataset.Na_cpu[nc],
+          dataset.virial_cpu[n + dataset.N * 4] / dataset.Na_cpu[nc],
+          dataset.virial_cpu[n + dataset.N * 5] / dataset.Na_cpu[nc],
+          dataset.avirial_ref_cpu[n],
+          dataset.avirial_ref_cpu[n + dataset.N],
+          dataset.avirial_ref_cpu[n + dataset.N * 2],
+          dataset.avirial_ref_cpu[n + dataset.N * 3],
+          dataset.avirial_ref_cpu[n + dataset.N * 4],
+          dataset.avirial_ref_cpu[n + dataset.N * 5]);
+        fprintf(
+          fid_stress,
+          "%g %g %g %g %g %g %g %g %g %g %g %g\n",
+          dataset.virial_cpu[n] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.virial_cpu[n + dataset.N] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.virial_cpu[n + dataset.N * 2] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.virial_cpu[n + dataset.N * 3] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.virial_cpu[n + dataset.N * 4] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.virial_cpu[n + dataset.N * 5] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.avirial_ref_cpu[n] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.avirial_ref_cpu[n + dataset.N] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.avirial_ref_cpu[n + dataset.N * 2] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.avirial_ref_cpu[n + dataset.N * 3] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.avirial_ref_cpu[n + dataset.N * 4] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION,
+          dataset.avirial_ref_cpu[n + dataset.N * 5] / dataset.structures[nc].volume * PRESSURE_UNIT_CONVERSION);
+      }
+    }
+  }


would it be possible to use/modify the existing output function, or alternatively, make a new function called output_atomic to clean this up a bit?

elindgren · 2025-02-21T15:38:09Z

src/main_nep/fitness.cu

+  if (!atomic) {
+    output(false, 3, fid_dipole, dataset.virial_cpu.data(), dataset.virial_ref_cpu.data(), dataset);
+  } else {
+    for (int nc = 0; nc < dataset.Nc; ++nc) {
+      int offset = dataset.Na_sum_cpu[nc];
+      for (int m = 0; m < dataset.structures[nc].num_atom; ++m) {
+        int n = offset + m;
+        fprintf(
+          fid_dipole,
+          "%g %g %g %g %g %g\n",
+          dataset.virial_cpu[n],
+          dataset.virial_cpu[n + dataset.N],
+          dataset.virial_cpu[n + dataset.N * 2],
+          dataset.avirial_ref_cpu[n],
+          dataset.avirial_ref_cpu[n + dataset.N],
+          dataset.avirial_ref_cpu[n + dataset.N * 2]);
+      }
+    }
+  }


same; try and reuse printing functionality if possible

elindgren · 2025-02-21T15:38:24Z

src/main_nep/fitness.cu

+  if (!atomic) {
+    output(false, 6, fid_polarizability, dataset.virial_cpu.data(), dataset.virial_ref_cpu.data(), dataset);
+  } else {
+    for (int nc = 0; nc < dataset.Nc; ++nc) {
+      int offset = dataset.Na_sum_cpu[nc];
+      for (int m = 0; m < dataset.structures[nc].num_atom; ++m) {
+        int n = offset + m;
+        fprintf(
+          fid_polarizability,
+          "%g %g %g %g %g %g %g %g %g %g %g %g\n",
+          dataset.virial_cpu[n],
+          dataset.virial_cpu[n + dataset.N],
+          dataset.virial_cpu[n + dataset.N * 2],
+          dataset.virial_cpu[n + dataset.N * 3],
+          dataset.virial_cpu[n + dataset.N * 4],
+          dataset.virial_cpu[n + dataset.N * 5],
+          dataset.avirial_ref_cpu[n],
+          dataset.avirial_ref_cpu[n + dataset.N],
+          dataset.avirial_ref_cpu[n + dataset.N * 2],
+          dataset.avirial_ref_cpu[n + dataset.N * 3],
+          dataset.avirial_ref_cpu[n + dataset.N * 4],
+          dataset.avirial_ref_cpu[n + dataset.N * 5]);
+      }
+    }
+  }


reuse output if possible

elindgren · 2025-02-21T15:39:54Z

src/main_nep/parameters.cu

@@ -172,7 +174,7 @@ void Parameters::calculate_parameters()
    lambda_e = lambda_f = 0.0f;
    enable_zbl = false;
    if (!is_lambda_v_set) {
-      lambda_v = 1.0f;
+      lambda_v = 1.0f; // by default, dipole or polarizability is fitted with global quantities


Is the same lambda_v used for atomic and global virials? You might have to double check the relative weight so that the default value of lambda_v still make sense for atomic virials.

I think it should be fine, since as far as I can tell lambda_v is multiplied by the per-atom virials in both cases already.

elindgren · 2025-02-21T15:42:26Z

src/main_nep/structure.cu

  std::ifstream& input,
  const Parameters& para,
  Structure& structure,
  std::string& xyz_filename,
-  int& line_number)
+  int& line_number,
+  int train_mode)


what does train_mode do? is it a keyword that indicates if it's atomic or global virials that are being used?

Yi-FanLi added 16 commits February 14, 2025 10:10

support using atomic virial to train potential, dipole, and polarizab…

b2cab1b

…ility

fix merge conflict

ac515b0

fall back to using only 3 modes

6320802

update fitness to support fitting atomic virial

ed44d68

correct reading avirial

ebc8702

copy avirial

da1c1cd

revert fitness

1f9779a

abandon lambda_av

9d11916

add option atomic_v

586af5d

use para.atomic_v to determine which rmse to use

eedc948

use atomic_v to determine whether to use global or atomic virial

1c899dc

Merge branch 'brucefan1983:master' into avirial

1a14318

remove avirial and avirial_cpu from dataset because they have already…

d0b2b1d

… been stored in virial

Merge remote-tracking branch 'refs/remotes/origin/avirial' into avirial

7a24cef

pass para.atomic_v to update_energy_force_virial, update_dipole, and …

f298611

…update_polarizability

add an example to training atomic dipole of water

b1ef48e

Yi-FanLi marked this pull request as ready for review February 18, 2025 05:59

Yi-FanLi added 6 commits February 18, 2025 01:15

add descriptions of atomic_v in doc

e290537

clarify the doc

2e5c652

do not use train_mode 4, 5, 6, and 7

4d738e4

correct atomic virial printing

ab6b232

correct the order of reading polarizability

53b8060

add atomic polarizability example

af5b508

brucefan1983 requested review from brucefan1983, erhart1, elindgren and tamaswells February 18, 2025 09:14

elindgren reviewed Feb 21, 2025

View reviewed changes

brucefan1983 marked this pull request as draft March 2, 2025 17:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support fitting atomic virial, dipole, and polarizability #907

Support fitting atomic virial, dipole, and polarizability #907

Yi-FanLi commented Feb 14, 2025 •

edited

Loading

elindgren left a comment

elindgren Feb 21, 2025

elindgren Feb 21, 2025

elindgren Feb 21, 2025

elindgren Feb 21, 2025

elindgren Feb 21, 2025

elindgren Feb 21, 2025

elindgren Feb 21, 2025

Support fitting atomic virial, dipole, and polarizability #907

Are you sure you want to change the base?

Support fitting atomic virial, dipole, and polarizability #907

Conversation

Yi-FanLi commented Feb 14, 2025 • edited Loading

elindgren left a comment

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

elindgren Feb 21, 2025

Choose a reason for hiding this comment

Yi-FanLi commented Feb 14, 2025 •

edited

Loading