Atom virial and stress calculation for virial calculation parallelism of Allegro in Lammps

[Hongyu-yu]

ParaStressForceOutput Support for atom virial and virial parallelism in Lammps.

Description

Another algorithm implementation for stress calculation can calculate atom virial and should help with parallelism of Allegro in Lammps.
PBC systems are used for tests.

Motivation and Context

Resolves: pari_allegro#9
Enable Allegro NPT simulations.

How Has This Been Tested?

Stress and virial results of ParaStressForceOutput and StressForceOutput are consistent.
Tests are carried on a PBC system with cell.
Just replace StressForceOutput with ParaStressForceOutput :)
See test_stress in nequip/tests/unit/model/test_nequip_model.py.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds or improves functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation improvement (updates to user guides, docstrings, or developer docs)

Checklist:

• My code follows the code style of this project and has been formatted using black.
• All new and existing tests passed, including on GPU (if relevant).
• I have added tests that cover my changes (if relevant).
• The option documentation (docs/options) has been updated with new or changed options.
• I have updated CHANGELOG.md.
• I have updated the documentation (if relevant).

[Linux-cpp-lisp]

👍

[Linux-cpp-lisp]

Is there a reason not to just use edge_virial.transpose(-1, -2) and avoid recomputing the einsum?

[Huni-ML]

Thanks for your suggestions and we can just do transpose.

[Linux-cpp-lisp]

This guarantees a symmetric virial?

[Huni-ML]

That's right.

[Linux-cpp-lisp]

Suggested change

# virial = grads[1]

[Linux-cpp-lisp]

This business is all unnecessary now, right? Since there is no derivative w.r.t. the cell, so the cell doesn't need to be made per-batch-entry

[Huni-ML]

Yes. Cell here is only used to get the stress and for the training with stress.

[Linux-cpp-lisp]

Suggested change

data[AtomicDataDict.CELL_KEY] = orig_cell

[Linux-cpp-lisp]

Would need to confirm correctness of sign convention here...

[Huni-ML]

Agree. I didn't check it and just inherit it from StressForceOutput .

[Linux-cpp-lisp]

Should also reset requires_grad opn the edge vectors

[Linux-cpp-lisp]

This fake data is never used, right?

[Huni-ML]

Right. This part is just inherited from molecules. Can just delete them.

[tgmaxson]

What is the status of this? If it is working, I do have an immediate use / test case for it and am very interested.

[Linux-cpp-lisp]

Hi @tgmaxson ,

Not sure about status of this implementation; our implementation using the existing StressForceOutput approach is available for Allegro with a prototype parallel integration into LAMMPS here: https://github.com/mir-group/pair_allegro/tree/stress

If you have further questions about that please feel free to open a thread in the pair_allegro repo or in one of the existing discussions across here and the allegro and pair_allegro repos.

[Hongyu-yu]

@Linux-cpp-lisp Hi, as you can see the tests all pass in python interface (give consistence result with StressOutput with cell) and Lammps interface (consistence result between python and Lammps with different ranks) and my own experience and checks, I am quite sure this PR is reliable and suitable for Allegro for NPT simulations in parallel so I don't think this PR and the PR of pair_allegro need more changes.

I have also tried StressOutput with Allegro in lammps in parallel but it gives inconsistent force and stress caused by pair_allegro interface.
So right now for my training and simulations, I use ParaStressForceOutput.

@tgmaxson You can have a try by using nequip and pair_allegro with my modification and train with ParaStressForceOutput, which I have tested a lot and been using it for a while and things goes well on my side.

The analysis of my implementations can be seen in page 11 of this preprint.

Feel free if you meet any problems.

[Linux-cpp-lisp]

Hi @Hongyu-yu ,

Ah great, thanks for the update!

I have also tried StressOutput with Allegro in lammps in parallel but it gives inconsistent force and stress caused by pair_allegro interface.

Interesting, was this with the stress branch of pair_allegro or no?

[Hongyu-yu]

Hi @Linux-cpp-lisp,

Right. I tested with the stress branch and it failed. I test it again just then (see here).

Hope that this PR helps.

[Linux-cpp-lisp]

@Hongyu-yu I see, thanks for clarifying and the quick responses!--- I'm hoping to find enough time to look through all this carefully this week.

[Hongyu-yu]

@Linux-cpp-lisp Thanks for your work:)

[Hongyu-yu]

Thanks for the explanations. I have no more questions regarding the virial/stress, except that I think it would also be nice to have testing results using LAMMPS (there are options to get the stress from both methods). Then it is safe to be used for applications requiring NPT or accessing virial distribution (as in tensile loading simulations).

As for heat current and heat transport, I am still not sure if the current implementation is totally correct. My feeling is the message passing can complicate the heat current. But I might find some LAMMPS experts to test it. We are interested in the high accuracy provided by Nequip but virial and heat current are crucial for our applications.

Thanks for sharing your thoughts. The implementation works fine for both Allegro and NequIP and for Allegro, which is strictly local, the atomic virial could be possible to be right and I think it's worthwhile to give it a try.
Hope it helps:)

[sirmarcel]

Hello everyone,

Sorry for barging in on this PR with some PR 😄, but I thought I'd highlight some developments that may be of interest to this discussion! I talked to Simon and Anders a while ago about these topics as well.

Regarding the heat flux, we wrote a small article a few weeks ago that confirms @brucefan1983's suspicion: Message passing makes the heat flux a bit more complicated, especially with AD. Luckily there's a fairly straightforward way to implement it anyway with AD. We also just released another article that goes into a bit more implementation detail (example implementation in JAX), and also discusses computing the stress (and atomic stresses) with AD.

It should be fairly straightforward to apply this heat flux to NequIP (or Allegro, where the "local" version is even faster). Let me know if I can help with that!

Cheers!

[Hongyu-yu]

Hello everyone,

Sorry for barging in on this PR with some PR 😄, but I thought I'd highlight some developments that may be of interest to this discussion! I talked to Simon and Anders a while ago about these topics as well.

Regarding the heat flux, we wrote a small article a few weeks ago that confirms @brucefan1983's suspicion: Message passing makes the heat flux a bit more complicated, especially with AD. Luckily there's a fairly straightforward way to implement it anyway with AD. We also just released another article that goes into a bit more implementation detail (example implementation in JAX), and also discusses computing the stress (and atomic stresses) with AD.

It should be fairly straightforward to apply this heat flux to NequIP (or Allegro, where the "local" version is even faster). Let me know if I can help with that!

Cheers!

Hi @sirmarcel ,
I have read your preprint today actually and it is a quite interesting:)
For my stress and atomic virial implementation, you may have a check the code or decription in this preprint which seems to be similar and suitable for heat flux. As I am not so familiar with heat flux, it will be helpful for an examination on the atomic virial. The rightness and accuracy of the system stress is checked in the tests. What I am not sure is about the atomic virial implementaion which is not easily checked. Thanks:)

[sirmarcel]

Hey,

The equations in the paper look indeed fine, but for heat flux, you need to be rather careful about indexing, as @brucefan1983 pointed out above: The heat flux is typically computed as sum_i W_i v_i where W_i is a virial, but the virial is sum_j r_ji \otimes d U_j / d r_ji and not, as one would maybe expect sum_j r_ij \otimes d U_i / d r_ij. (Modulo sign conventions, etc.; this also only holds for strictly local potentials, i.e. no message passing)

In other words, the index of the atomic energy shouldn't be the same index as the one in the velocity. I'm not sure what happens if you do the averaging you describe in the text after eq. 9. Is there any reason you do it this way?

[Hongyu-yu]

In other words, the index of the atomic energy shouldn't be the same index as the one in the velocity. I'm not sure what happens if you do the averaging you describe in the text after eq. 9. Is there any reason you do it this way?

Hi @sirmarcel ,
Thanks for the check. Here I first calculated the edge_virial r_ij \otimes d U_i / d r_ij and them accumulated and distributed them into two nodes equally in code. As you suggested, these code should be deleted, right?

[sirmarcel]

Sorry, I'm not so familiar with the internals of the nequip code -- what you need to do to compute the virial for atom j is to sum up all terms where the edge contains that index, and not the potential energy: d U_i / d r_ij not d U_j / d r_ij. I think that's what

 scatter(
            edge_virial,
            data[AtomicDataDict.EDGE_INDEX_KEY][1],
            dim=0,
            reduce="sum",
            dim_size=len(pos),
        )

in your code is doing, but I'm not sure, as I don't really know that the index arrays are, etc. In all cases you should get the same total stress tensor for the system.

[brucefan1983]

Hello, sorry for disturbing again. Thanks @sirmarcel for confirming that heat current for message-passing construction is a little more complicated than "local" machine-learned potentials. I also read your preprints and they look very nice.

I think the current parallel version of virial by @Hongyu-yu should be correct as he confirmed there are consistency between two approaches (system virial and atomic virial). I suggest merging this nice work to master branch. This feature is crucial for efficient NPT simulations and alike in LAMMPS. I have seen this PR to be here for a long time. Why not move on?

Heat current is more involved and can be left for a future work. One can put a warning message though.

[Hongyu-yu]

Hi @sirmarcel ,
In my code, edge_virial is $r_{ji}\otimes d U/d r_{ji}=r_{ji}\otimes d U_i/d r_{ji} + r_{ji}\otimes d U_j/d r_{ji} + \sum_{k!=i,j} r_{ji}\otimes d U_k/d r_{ji}$ with two atomic energy contribution from the two atoms on the edges and the neighbor atoms around the bond. However, the detailed contribution of three parts can not be distinguished. To be noted, the atomic energy itself is an assumption from Behler and Parinello which can not be strictly defined or measured from the experiments.
For atom_viral as you suggested, should be $\sum_j r_{ji} \otimes d U_j/d r_{ji}$, which seems to hard to built with the edge_virial as they are mixed together. So I mix them together with sum_i and sum_j.
Maybe use the atomic energy directly in the autodifferentiation to get atomic virial with $\sum_j r_{ji} \otimes d U_j/d r_{ji}$ is better? What do you think?

[Hongyu-yu]

Hello, sorry for disturbing again. Thanks @sirmarcel for confirming that heat current for message-passing construction is a little more complicated than "local" machine-learned potentials. I also read your preprints and they look very nice.

I think the current parallel version of virial by @Hongyu-yu should be correct as he confirmed there are consistency between two approaches (system virial and atomic virial). I suggest merging this nice work to master branch. This feature is crucial for efficient NPT simulations and alike in LAMMPS. I have seen this PR to be here for a long time. Why not move on?

Heat current is more involved and can be left for a future work. One can put a warning message though.

Hi @brucefan1983,
thanks for the comments:) The system virial with my code works pretty fine while the atomic virial should be tested more as there is not a value to compare with for the check right now.
Accutually I am also quite worried about the correctness and the influence of the stress implementation in the stress branch of nequip and pair_allegro. The combination of nequip and pair_nequip works fine as the cell information is taken into the LAMMPS shown in this test. However, it works wrong actually with the pair-allegro shown in the tests as the cell information is missed, which influence the auto differentiation. It will surely give unreliable results to users.
This implementation could be a better choice with strict theory and careful numerial checks. The cell information can be negelected and suitable for the parallel computation.
As @brucefan1983 mentioned, the atomic virial should be tested more and may need some updates while it doesn't influence the system virial acutually.

[enikidis]

Hello @Hongyu-yu and thanks a lot for you implimetation of per atom virial stress calculation. I am using your para_stress fork of nequip and your pair_allegro stress to train on the Copper xyz dataset file provided in the mir group repo.
The file looks like:

8
Lattice="0.0 3.61 3.61 3.61 0.0 3.61 3.61 3.61 0.0" Properties=species:S:1:pos:R:3:forces:R:3:energies:R:1 energy=-0.03845436888363629 free_energy=-0.03845436888363629 pbc="T T T"
Cu       0.01764052       0.00400157       0.00978738       0.00032332       0.01084934      -0.02260498      -0.00535298
Cu       1.82740893       1.82367558      -0.00977278      -0.08662083      -0.06473697       0.13650100      -0.00475540
Cu       1.81450088      -0.00151357       1.80396781       0.03243655      -0.00841086      -0.09314117      -0.00459684
Cu       3.61410599       1.80644044       1.81954274       0.03938003       0.01648442      -0.12443480      -0.00473976
Cu       0.00761038       1.80621675       1.80943863       0.01666194       0.06380369      -0.15080920      -0.00441283
Cu       1.80833674       3.62494079       1.80294842       0.00467744      -0.04715308      -0.01894000      -0.00560490
Cu       1.80813068       1.79645904       3.58447010       0.02970473       0.08976916       0.19657489      -0.00378700
Cu       3.61653619       3.61864436       3.60257835      -0.03656318      -0.06060571       0.07685426      -0.00520466

I am trainig the potential using the following yaml with your new Allegro model ParaStressForceOutput. And the training is finished successfully and I deploy the pth file.
I am using the following lammps script to calculate both the global stresses and the per_atom stresses for 4 copper atoms.

#read_restart restart1.Sl
units		metal
atom_style	atomic
newton on
thermo 1




# Atom Definition
read_data structure.data

pair_style	allegro3232
pair_coeff	* * CU-para-deployed.pth Cu
mass            1 63.546

neighbor	1.0 bin
neigh_modify    delay 0 every 1 check no

fix		1 all nve

timestep	0.001


compute global_stress all pressure NULL virial  # NULL means without temperature contribution
thermo_style custom step c_global_stress[*]
compute per_atom_virial all centroid/stress/atom NULL virial
dump	1 all custom 1 dump_peratom_stress.atom x y z c_per_atom_virial[*]


run 2

What I get as an output for the global

tep     c_global_stress[1] c_global_stress[2] c_global_stress[3] c_global_stress[4] c_global_stress[5] c_global_stress[6]
         0   129947.2       129947.16      129947.14     -0.01351072    -0.012745496   -0.017340643  
         1   129947.25      129947.19      129947.19      0.0084408317   0.0078877044   0.0081367217 
         2   129947.25      129947.23      129947.14      0.0075560226   0.0064573176   0.01289513          

and the per atom stress are always zero:

ITEM: TIMESTEP
0
ITEM: NUMBER OF ATOMS
4
ITEM: BOX BOUNDS pp pp pp
0.0000000000000000e+00 3.5970000000000000e+00
0.0000000000000000e+00 3.5970000000000000e+00
0.0000000000000000e+00 3.5970000000000000e+00
ITEM: ATOMS x y z c_per_atom_virial[1] c_per_atom_virial[2] c_per_atom_virial[3] c_per_atom_virial[4] c_per_atom_virial[5] c_per_atom_virial[6] c_per_atom_virial[7] c_per_atom_virial[8] c_per_atom_virial[9]
0 0 0 -0 -0 -0 -0 -0 -0 -0 -0 -0
0 1.7985 1.7985 -0 -0 -0 -0 -0 -0 -0 -0 -0
1.7985 0 1.7985 -0 -0 -0 -0 -0 -0 -0 -0 -0
1.7985 1.7985 0 -0 -0 -0 -0 -0 -0 -0 -0 -0
ITEM: TIMESTEP
1
ITEM: NUMBER OF ATOMS
4
ITEM: BOX BOUNDS pp pp pp
0.0000000000000000e+00 3.5970000000000000e+00
0.0000000000000000e+00 3.5970000000000000e+00
0.0000000000000000e+00 3.5970000000000000e+00
ITEM: ATOMS x y z c_per_atom_virial[1] c_per_atom_virial[2] c_per_atom_virial[3] c_per_atom_virial[4] c_per_atom_virial[5] c_per_atom_virial[6] c_per_atom_virial[7] c_per_atom_virial[8] c_per_atom_virial[9]
3.597 3.597 3.597 -0 -0 -0 -0 -0 -0 -0 -0 -0
7.64676e-13 1.7985 1.7985 -0 -0 -0 -0 -0 -0 -0 -0 -0
1.7985 3.597 1.7985 -0 -0 -0 -0 -0 -0 -0 -0 -0
1.7985 1.7985 3.597 -0 -0 -0 -0 -0 -0 -0 -0 -0

Probably I haven't understood something correctly about the functionality in LAMMPS.
Any clues why this happens? Is it possible that the dataset is too small, thus the potential is not very good quality?

[Hongyu-yu]

Hi，@enikidis
Thanks for your interest!
With some trials and check of source code, I found that it is caused by the vflag_atom in Lammps. In your script, the vflag_atom is False, which should be activated to be true. For a check of the atomic virial, you can mannually modify if (vflag_atom) to if (vflag) in line 495 of pair_allegro.cpp. Recompile the Lammps and it should give results instead of 0.

[pobo95]

Hi @Hongyu-yu,
Thank you so much for your code bucause it was a featur i needed.

But I tried to compile your nequip and pair_allegro in Lammps with Pythoch version 1.10.2

I got this error massage

lammps-stable/src/pair_allegro.cpp(512): internal error: null pointer }

Can you know me how to fix it?

Thanks again!

[Hongyu-yu]

Hi, @pyungjinpark
Thanks for your interest! Have you made a change on the code? It seems like an error caused by the unconsistent bracket, which is correct for the source code.

[pobo95]

Thank you for your quick reply @Hongyu-yu !

I didn't change your sorce file, I will try it agin.

And I tried 2000K NPT simulation in silicon with pair_nequip based on your code (you can check it in my repositories https://github.com/pyungjinpark/pair_nequip).

It give me atomic press well but when I sum it, It is different with total press

Step	Temp	Press	v_press	c_p[1]	c_p[2]	c_p[3]	Volume
0	2000	29552.29	-2546.95	92884774	1.19E+08	1.01E+08	40879.62
1000	1906.324	6107.532	19090.08	-7.7E+08	-7.8E+08	-8.6E+08	42133.72
2000	1937.251	-3211.19	27518.95	-1.1E+09	-1.2E+09	-1.2E+09	43243.92
3000	2039.873	-7967.44	33364.15	-1.5E+09	-1.4E+09	-1.5E+09	43261.44
4000	1997.884	-4648.99	30553.3	-1.2E+09	-1.3E+09	-1.3E+09	42736.52
5000	2000.029	-354.636	25967.71	-1.1E+09	-1.1E+09	-1.1E+09	41959.05
6000	2012.074	-1919.34	30352.47	-1.3E+09	-1.2E+09	-1.3E+09	41261.93
7000	2005.202	3314.31	23978.89	-9.8E+08	-9.5E+08	-1E+09	40850.05
8000	2043.332	1910.847	25981.38	-1.2E+09	-1E+09	-9.6E+08	40711.35
9000	1982.122	-608.586	26889.21	-1.1E+09	-1.2E+09	-1E+09	40771.17
10000	1978.882	2325.762	23693.38	-1E+09	-1E+09	-8.9E+08	40952.3
11000	2039.672	670.9164	27113.22	-1.1E+09	-1.1E+09	-1.1E+09	41158.75
12000	2030.595	-578.722	28673.16	-1.3E+09	-1.1E+09	-1.2E+09	41315.22
13000	1998.566	-901.163	28050.84	-1.2E+09	-1.2E+09	-1.1E+09	41378.63
14000	1962.711	-1298.09	28580.57	-1.1E+09	-1.2E+09	-1.2E+09	41339.99
15000	1969.174	-3286.15	29650.17	-1.2E+09	-1.2E+09	-1.2E+09	41248.88
16000	2052.363	-1170.49	28327.81	-1.3E+09	-1.2E+09	-1E+09	41165.25
17000	2045.011	1996.934	26907.24	-1.2E+09	-1E+09	-1E+09	41111.18
18000	2000.428	-244.389	28071.88	-1.2E+09	-1E+09	-1.2E+09	41109.37
19000	2004.759	245.2562	26635.44	-1.2E+09	-1E+09	-1E+09	41130.73
20000	1975.086	-1642.7	27732.93	-1.2E+09	-1.1E+09	-1.1E+09	41166.55

where Press is total press and v_press is sum of atomic press

$${-1 \over 3*volume} \sum_{i=1}^N (p_ixx +p_iyy + p_izz) $$

When I tested same thing in QUIP, it gave same total press and v_press

Do you have any clue?

[tgmaxson]

just a quick thought, does the pressure in this test you are doing include the ideal gas contribution from velocity? this may result in deviations so make sure you are calculating things consistently (which is the normal error I have when moving from LAMMPS to analysis).

[pobo95]

@tgmaxson

Thank you for your comment sir!
When we calcualte the stress tensor I used this formula in lammps
$$S_{ab} = -mv_{a}v_{b} -W_{ab} $$

So Yes, I include the ideal gas contribution from velocity (kinetic energy)
In my experience, when we do not get (average P)=0 virial was the problem.

And when I draw Step vs press and Step vs v_press graph

It looks like has sign problem.
Maybe I shold fix here

https://github.com/pyungjinpark/pair_nequip/blob/848b2c381d46a3bdcfd17c04e6643e5ab23d514a/pair_nequip.cpp#L494-L510

[Hongyu-yu]

Hi @pyungjinpark ,
How about the results after change the sign for atomic virial? It could be wrong sign in the source code.

[pobo95]

@Hongyu-yu

It is solved!

Step	Press	v_press	c_p[1]	c_p[2]	c_p[3]	Volume
0	29552.29	29552.28	-1.2E+09	-1.2E+09	-1.2E+09	40879.62
1000	6142.151	6142.15	-2.9E+08	-2.1E+08	-2.7E+08	42133.33
2000	-10770	-10770	4.89E+08	4.69E+08	4.39E+08	43240.7
3000	-4186.88	-4186.88	50649014	1.83E+08	3.1E+08	43254.46
4000	-1216.02	-1216.02	82237257	-2.1E+07	94810384	42732.64
5000	-1562.84	-1562.84	1.13E+08	66590482	16858399	41966.78
6000	-1046.23	-1046.23	14484347	48763754	66306025	41276.4
7000	2416.496	2416.498	-7.3E+07	-1.5E+08	-7E+07	40844.99
8000	4955.038	4955.036	-2.4E+08	-2.2E+08	-1.4E+08	40687.91
9000	1048.96	1048.965	-1.7E+08	28420884	12695171	40739.96
10000	4320.356	4320.351	-2.3E+08	-9.1E+07	-2.1E+08	40909.87
11000	1155.951	1155.945	-3.2E+07	-4.7E+07	-6.4E+07	41098.12
12000	-438.924	-438.922	61134988	-2.2E+07	14884850	41248.56
13000	-2321.57	-2321.56	66696030	1.45E+08	76578624	41324.62
14000	-52.6433	-52.6388	-6.8E+07	-5.4E+07	1.28E+08	41317.44
15000	582.6992	582.7009	-1E+08	-1.6E+07	47562333	41249.31
16000	1820.461	1820.461	-1.9E+08	-2.6E+07	-1.1E+07	41160.78
17000	-117.581	-117.576	-3.3E+07	-2684659	49943344	41121.85
18000	-743.572	-743.574	-1.2E+07	93638421	10187999	41109.2
19000	-1686.68	-1686.68	1.11E+08	53471306	43924974	41125.48
20000	2042.577	2042.578	-9.5E+07	-5.3E+07	-1E+08	41144.51

It shows some error (not exactly same) but works well when I fix the sign.
You can check what I changed
https://github.com/pyungjinpark/pair_nequip/blob/a1e88d0f4c06542fe5084947dbd25557b89f54e9/pair_nequip.cpp#L494-L505

[Hongyu-yu]

@pyungjinpark Looks great! How about the results that only changing the sign of atomic virial? The atomic virial is not required to be symmetric.

[pobo95]

Hi !@Hongyu-yu

1. It gives the same result because we use only xx, yy, zz atomic virials to get total press.
2. And I initially coded it not to be symmetric (to get 9 atomic virial)
   but Lammps gave only 6(only xx, yy, zz, xy, xz, yz except yx, zx, zy).

The definition of virial in lammps is

$$ W^{pair} = \sum_{i=1}^N \sum_{j>i}^N \vec{r}_{ij} \otimes \vec{F}_ij $$

where
r is the relative position of atom i in relation to atom j
F is the force action upon atom i due to interaction with atom j

In case of heat flux vector

$$ \vec{J} = {1 \over V} \left\langle \sum_{i=1} \vec{v}_i e_i + \vec{Q} \right\rangle $$

where
Q is virial contribution

$$ \vec{Q}^{pair} = - \sum{\sigma}_{i}^{pair} \cdot \vec{v}_i $$

Maybe Lammps assumes it is symmetric.

If I understand corretly you used edge data for virial then It might give symmetric information, isn't it?

4. I check another ML potental (QUIP), it used same way with me.

[Hongyu-yu]

Hi @pyungjinpark ,
Thanks for the detailed explanation! Can you give the exact lines of the implementation of QUIP?
As @brucefan1983 mentioned above,

Particularly, it has been emphasized that the per-atom virial cannot be assumed to be symmetric. If you want to get the correct heat current in LAMMPS, the interface must provide 9 components as detailed in the following page:
https://docs.lammps.org/compute_stress_atom.html
Atomic virial should have nine components. If you only change the sign without symmetrization, are the results wrong?

[pobo95]

Here is the exact lines of QUIP @Hongyu-yu

if (vflag_atom) {
int iatom = 0;
for (ii = 0; ii < ntotal; ii++) {
vatom[ii][0] += quip_local_virial[iatom + 0];
vatom[ii][1] += quip_local_virial[iatom + 4];
vatom[ii][2] += quip_local_virial[iatom + 8];
vatom[ii][3] += (quip_local_virial[iatom + 3] + quip_local_virial[iatom + 1]) * 0.5;
vatom[ii][4] += (quip_local_virial[iatom + 2] + quip_local_virial[iatom + 6]) * 0.5;
vatom[ii][5] += (quip_local_virial[iatom + 5] + quip_local_virial[iatom + 7]) * 0.5;
iatom += 9;
}
}

you can check it in lammps/src/ML-QUIP/pair_quip.cpp

@brucefan1983 is right.

I only considered pair-atomic like virial, I will change the code !
It will give better results for heat flux.

[Hongyu-yu]

Hi, @pyungjinpark
Thanks! In my code, it uses cvatom instead of vatom for atomic virial. You may check the difference in https://docs.lammps.org/compute_stress_atom.html.
I have fixed the sign in my pair_allegro. Hope it works fine for you:)