Computing/Predicting per-atom virials using the allegro model

[eshaanraval]

I'm trying to implement the HNEMD approach (Paper) for determining thermal conductivity in torch_sim, using MLIPs like Allegro and so on. I'm new to training ML potentials and looking for guidance on how can I train an allegro model to predict per-atom virials along with all the other stuff it predicts in the pre-trained model like mir-group/Allegro-OAM-L.

For the approach I need to compute per-atom virial tensors based on a certain formula:

$$W_i = \sum_{i \ne j}{\frac{\partial U_j}{\partial \mathbf{r}_{ji}} \otimes \mathbf{r}_{ij}}$$

where $\vec{r_{ji}}$ $\equiv$ $\vec{r_i} - \vec{r_j}$ is the displacement vector between atoms $i$ and $j$, and $U$ is the total potential energy.

Is it possible to train my own model to predict the per_atom virials along with the other quantities

Like Would it be possible to:

• Train an Allegro model to predict per-atom virial tensors as direct outputs?
• Use the virial tensor formula above as a training target during model training?
• How difficult would it be to implement this?
• How would/should one go about implementing this and training their own model?

Any clarification or guidance on training a custom model for direct virial prediction, would be greatly appreciated!