Help on virial calculation using virial_fdotr_compute()

Dear lammps developers,

We always thank you for your invaluable help!

My colleague is writing a pair style for a GPU version of a machine learning potential using the JAX library. For a CPU version, after assigning energies and forces, the calculation of stress was done using
if (vflag_fdotr) virial_fdotr_compute();
at the end of the ‘compute’ function and it worked to give correct values of stress (So the corrected version looks like this: LAMMPS/pair_MLIP_4.cpp · main · Alexander Shapeev / interface-lammps-mlip-4 · GitLab).

But for this GPU version, the same thing was done as follows.

  // 4. PROCESS RESULTS
  // JAX returns forces for all nlocal atoms. We apply them only to the ones in ilist.
  for (int ii = 0; ii < inum; ii++) {
    int i = ilist[ii];
    f[i][0] += persistent_forces_array[i][0];
    f[i][1] += persistent_forces_array[i][1];
    f[i][2] += persistent_forces_array[i][2];
  }

  if (eflag_global) eng_vdwl += total_energy;
  if (vflag_fdotr) virial_fdotr_compute();
  total_calls++;
}

And the energy and forces are correct as expected, but the stress values are different. Maybe this virial_fdot_compute() function is more than simple summation of f dot r?

The machine learning potential belongs to the category of manybody potential, because to get the site (atomic) energy and the atomic force, only the coordinates within the cutoff radius are used. There is no added Coulomb interaction, no extra communication out of the cutoff radius. In the end the site energies are summed to give the total energy.

Thank you very much!
Sincerely,
Jong

The Pair::virial_fdotr_compute() function is only valid to use when Force::newton_pair is 1 which means that the forces on ghost atoms have not yet been summed to their corresponding local atoms during the reverse communication. So, your CPU version either has to disallow using Force::newton_pair == 0 or call the suitable Pair::ev_tally*() function and/or Pair::v_tally*() function as needed during the force computation.

The same would apply to the GPU version. I suspect that one is not keeping the forces on the ghost atoms in the same way the CPU version does and that would explain the difference.

Please also note: ML-IAP with JAX by wcwitt · Pull Request #4691 · lammps/lammps · GitHub