Virial and Lammps interface about pair_allegro

[Linux-cpp-lisp]

Hi @Huni-ML ,

Wow! 😆 Very nice!

I was waiting for it to be a little further along before I sent the link here, but we very recently have (most) of a finished stress implementation on the stress branch: https://github.com/mir-group/pair_allegro/tree/stress. We are, however, using a different approach than your linked PR: we are using the same dE_tot / d displacement approach as used in pair_nequip for a single node with LAMMPS automatic sum reduction over ranks. We've been verifying the correctness of this against single node stresses computed without LAMMPS, see the comprehensive unit tests here: https://github.com/mir-group/pair_allegro/blob/stress/tests/test_python_repro.py#L277-L300. It could be good to verify the correctness of your approach across different numbers of parallel ranks using the linked unit tests.

Please feel free also to do your LAMMPS tests using our stress branch, and let us know if you see any issues, or inconsistencies between the two approaches.

Do you have any thoughts on these two different approaches to computing the stress? Are they equivalent? Are there important advantages/disadvantages of one or the other?

Thank you for your efforts to contribute back into our codebase!

[Linux-cpp-lisp]

Hi @Huni-ML ,

Regarding NPT / stress / virial support in pair_allegro, this is something we are actively working on right now and hope to have out in beta soon--- I'll keep you posted. Depending on the size of your system, you can temporarily use your Allegro models with pair_nequip, which supports NPT, on one GPU.

Regarding your other questions, @anjohan ?

[anjohan]

Hi @Huni-ML ,

The Allegro model simply does not output stress, so this won't work. Alby is working on implementing it. Unlike NequIP, the Allegro model does not get a cell with respect to which it can take a derivative.

For the time being, you can run NPT on a single GPU using pair_nequip with the Allegro model.

The fdotr approach is, to my understanding, only used for non-periodic systems. Otherwise, you need to operate on relative distances and pairwise/partial forces.

As for eng_vdwl, it is summed as for any other potential.

  eng_vdwl = 0.0;
#pragma omp parallel for reduction(+:eng_vdwl)
  for(int ii = 0; ii < ntotal; ii++){
    int i = ilist[ii];

    f[i][0] = forces[i][0];
    f[i][1] = forces[i][1];
    f[i][2] = forces[i][2];
    if (eflag_atom && ii < inum) eatom[i] = atomic_energies[i][0];
    if(ii < inum) eng_vdwl += atomic_energies[i][0];
  }
}

[Huni-ML]

Thanks for your reply! I will use pair-nequip first.
For the stress calculation without cell, following lines may help:

        Z = torch.autograd.grad(
            energy,
            rij_vectors,
            create_graph=self.training,
            retain_graph=self.training,
        )

        result["stress_volume"] = torch.matmul(
            Z[0].t(), rij_vectors
        )  # Only consider one graph here

In my own package, it gives similar results with the cell-derivative strain. Seems to be suitable for Allegro.

[Linux-cpp-lisp]

@Huni-ML to clarify, pair_allegro in LAMMPS doesn't yet support stress/virial. The Allegro energy model itself can be used with StressForceOutput from nequip just like the NequIP energy model; this is what Anders is referring to when he says "you can run NPT on a single GPU using pair_nequip with the Allegro model." This combination also allows you to train an Allegro model using stress / virial data.

[Huni-ML]

@Huni-ML to clarify, pair_allegro in LAMMPS doesn't yet support stress/virial. The Allegro energy model itself can be used with StressForceOutput from nequip just like the NequIP energy model; this is what Anders is referring to when he says "you can run NPT on a single GPU using pair_nequip with the Allegro model." This combination also allows you to train an Allegro model using stress / virial data.

Thanks for the clarification and that's what I'm doing. Thank you for the reminder! I have used StressForceOutput but didn't use stress data for training as my supervisor suggests that stress from VASP is not quite reliable. Maybe I will try to use it and compare the results.

Best Regards!

[anjohan]

@Huni-ML We usually train on stress and find that it helps with getting the correct lattice parameter. You may, however, need to re-check the DFT convergence, as stress can be harder to converge than energy and forces.

[Huni-ML]

@anjohan I only collect calculations that reach DFT convergence so I think it works. Thanks!

[Huni-ML]

@Linux-cpp-lisp @anjohan Hi, I submit a pull request that could be used for Allegro virial parallelism calculation just then. It's so urgent for me so I try it myself. With the atom virial without usage of cell, it should be suitable for allegro.
So late in my time so more tests in Lammps tomorrow. Hope it works!

[Huni-ML]

Hi @Linux-cpp-lisp ,
Glad to be able to contribute back to your great work. Actually, all my recent works are based on Allgero and it's my pleasure to make it better.
In my experiments, StressForceOutput with pair_allegro gives wrong results even on the forces while works well with pair_nequip. So strange for me and I think it's caused by the cell input. pair_allegro without cell as input may cause incorrectness in

            data[AtomicDataDict.CELL_KEY] = cell + torch.bmm(
                cell, symmetric_displacement
            )

I'm not sure about the correctness of this approach though.
About the comparison of two approaches, my approach focuses on the edge and your approach focus on the nodes. The edge_vector is directly given by edge_index and positions in python and Lammps interface at the same time without cell. Node based approaches actually means giving a virtual force on the cell and calculate its response. It works right on pari-nequip but not sure on pair-allegro without a cell input. That means some gradients on the edge will be lost in Lammps interface as in with_edge_vectors

                edge_vec = edge_vec + torch.einsum(
                    "ni,nij->nj", edge_cell_shift, cell[batch[edge_index[0]]]
                )

In a word, 'edge_vec' calculated in pari-allegro interface has no gradients on cell while it should have just as in python and pari-nequip interface, which makes pari-allegro wrong in stress I think.
In general, they are equivalent as the edge_vector is derived from positions. Key here is to avoid the cell as the interface.
In the edge approach, I just calculate gradients over 'edge_vec' so this problem is solved. I will do some tests today and hope that all is well :)

[Huni-ML]

Hi @Linux-cpp-lisp
With some small changes in my nequip fork and my pair-allegro fork, all Lammps tests passed in this actions.
Happy to see it works:)
Strict derivation of virial is shown below based on the edge vector. I will update more theoretical analysis on it in this preprint. If ParaStressForceOutput is used, this preprint can be considered to cite:). Thanks!

[Linux-cpp-lisp]

Great @Huni-ML !!

Just to fill in the gaps for myself in your derivation: in general,

$$
\frac{\partial E_i}{\partial \vec{r}j} = \sum{k \neq j}{ \frac{\partial E_i}{\partial \vec{r}{jk}} + \frac{\partial E_i}{\partial \vec{r}{kj} } }.
$$

If $j \neq i$,

$$
\frac{\partial E_i}{\partial \vec{r}j} = \frac{\partial E_i}{\partial \vec{r}{ij} }
$$

due to Allegro's strict locality which implies that

$$ \forall i,j : \frac{\partial E_i}{\partial \vec{r}_{kj} } \neq 0 \iff k=i. $$

If $j = i$,

$$
\frac{\partial E_i}{\partial \vec{r}i} = \sum{k \neq i}{ \frac{\partial E_i}{\partial \vec{r}{ik}} + \frac{\partial E_i}{\partial \vec{r}{ki} } } = \sum_{k \neq i}{ \frac{\partial E_i}{\partial \vec{r}_{ik}} }
$$

once again by strict locality of Allegro.

So your derivation follows by:

1. Definition of virial
2. Expansion of total energy into site energies & definition of force
3. Above substitutions from strict locality
4. Rename indexes
5. Since

$$ \vec{r}_{ij} = \vec{r}_j - \vec{r}_i $$

isn't the minus sign wrong? And I'm assuming

$$ \sum_{i \neq j}{} $$

in this line only refers to

$$ \sum_{i}{\sum_{j \neq i}{}}. $$

Does this all seem correct to you?

[Huni-ML]

Hi @Linux-cpp-lisp ,
Your analysis is correct for me and I think you have already fully understood my method:)
Only some details differ.
The first equation should only contain the first term. rjk and rkj are no different but a sign.
The definition of rij is not consistent and I will take yours.

The minus sign below is carefully checked and should exists so now it should give the correct sign for Lammps. I will do some NPT simulations to ensure its corectness with some simple examples in the future.

virial = torch.neg(virial)

I make a throughout analysis of this edge-vector-based atomic virial method in the new version of this preprint in reseach gate and you can check more details there in page 11.
Glad that I can contribute to Allegro package to make it better:)
By the way, would you mind adding this preprint as the citation suggestion in References and citing part of Allegro package if ParaStressForceOutput is used by user? Thanks!

Best regards.

[hyuntae-cho]

Hi everyone, I'm recently looking for a way to implementing virials.
from what I understand, LAMMPS use ghost atoms for periodic boundary conditions and therefore
if (vflag_fdotr) virial_fdotr_compute();
can be used. I've checked agreement with if (vflag_fdotr) virial_fdotr_compute(); and StressForceOutput.

Thanks for reading this comment!