Atomic energy

[tjgiese]

I followed the Si tutorial https://colab.research.google.com/drive/1yq2UwnET4loJYg_Fptt9kpklVaZvoHnq#scrollTo=bsF-uTLjhxoy to generate si-deployed.pth. However, I cannot get allegro to produce a result when the system consists of a single Si atom. It also fails when the system consists of 2 Si atoms separated by a distance larger than r_max. Based on equations 5 and 6 in https://doi.org/10.1038/s41467-023-36329-y I would expect a single atom to return the value of the atomic shift parameter. Similarly, I would expect a well-separated dimer to return the sum of atomic shift parameters. It doesn't seem possible to calculate relative energies using the supermolecular approach [dE(AB)=E(AB)-E(A)-E(B)] when one of the monomers is an atomic ion, for example. I attached a small python script that reads a TorchScript model prepared with nequip-deploy and an ASE extended XYZ file. The script evaluates the model with each structure and outputs a new xyz file with the MLP energies and forces. I also attached the error message I receive when I run it on a system containing a single atom. (The script merely follows what pair_allegro.cpp does -- although I ignore PBCs).

ERR.txt
run_nequip.py.txt

[Linux-cpp-lisp]

Hi @tjgiese ,

Can you retry this test training a new model with the very latest develop branch, and see if it is resolved? I think it should be.

You can also use (and or slightly modify) the plot_dimers.py script to make these kinds of tests: https://github.com/mir-group/nequip/blob/develop/examples/plot_dimers.py

Regarding the single isolated atom, the easiest test is probably just to use plot_dimers.py with a large --r-max and see what happens in the limit when they are further apart than the model's cutoff (you should get the sum of the desired / set isolated atom energies for the two relevant atoms).

[tjgiese]

Hi @Linux-cpp-lisp ,

I'm also having difficulty with the develop branch. I retried with the latest nequip/develop branch. I reperformed the Si tutorial and used plot_dimers.py, as suggested. The tutorial trains to r_max=5.

• If I run plot_dimers.py --r-min=4.9 --r-max=5.2 --n-samples 2 si-deployed.pth, then it will display (see note below) the energy at 4.9 and 5.2, as expected.
• If I run plot_dimers.py --r-min=5.1 --r-max=5.2 --n-samples 2 si-deployed.pth, then plot_dimers.py segmentation faults.

To get the --r-min=4.9 case to work, I had to make 2 small changes to plot_dimers.py. These changes do not solve the segmentation fault described in the 2nd case.

The line reading
rs_one = torch.linspace(args.r_min, model_r_max, 500, device=args.device)
should instead read
rs_one = torch.linspace(args.r_min, model_r_max, N_sample, device=args.device)
and the line
ax = axs[i]
should read

    try:
        ax = axs[i]
    except:
        ax = axs # scalar, not array because there's only 1 type

Thanks, -Tim

[tjgiese]

This appears to be caused by the way that tensors are being reshaped. Allegro and nequip will often reshape with -1, causing several of the dimensions to collapse; however, pytorch complains about this usage when one of the dimension sizes is zero. In the first case (discussed in the previous comment), plot_dimers.py sends all the calculations to pytorch in a batch and -- because at least one of the structures in the batch has an edge -- the tensor dimensions are at least 1. In the second case, all the structures in the batch have no edges. The second case is more relevant when integrating allegro into MD packages because one typically only evaluates one structure at a time.

The changes needed to calculate a well-separated dimer are quite small and isolated within the allegro source code. A pull request can be found here: #84

Convincing pytorch to calculate the energy of an isolated atom is a little more cumbersome because pytorch's autograd will complain that there are no variables requiring gradients. The changes needed to get isolated atoms working is isolated within the nequip code. A pull request can be found here: mir-group/nequip#426