How to handle with small distance among atoms (GPU)?

[theAfish]

Hi all,

Recently I have been trying to play with the source code of deepmd-kit (v2.2.8.dev0+g839f4fe7.d20231129, TensorFlow version:2.8.0). I would like to use some "pseudo-atoms" to prob the scalar field near atoms. In this case, there will be some pseudo-atoms in a very short distance (<=0.04 or even =0 Å). It seems that when the distance is very small, there will be nan values while training on GPU, even if I have dropped the overlapped pseudo-atoms with atoms. I regarded this as the issue of 1/r terms. So, I was wondering whether I could first add a small value epsilon into the codes to test whether the implementation works.

I have tried to add an eps=0.1 to lib/src/gpu/prod_env_mat.cu, lib/src/env_mat.cc, and lib/include/ComputeDescriptor.h. For example in prod_env_mat.cu:

Since I would like to first test whether the training will work properly or not, I have not considered the correction of derivatives yet.

However, there are still nan values after training for 300 steps (min nbor dist: 0.04066720115198197). So here I would like to ask:

1. Is there any simple solution to deal with the short-distance problem?
2. Could anyone tell me where to change the codes if possible?

I will be very appreciative of any help you can provide! Thank you very much!