Enhancement: Heat Flux Function

[stefanbringuier]

Summary

• Add calc_heat_flux function to compute heat flux vectors
• Supports both virial and convective terms
  • Can request virial term only
• Optional centroid stress tensor form
• Supports batch processing

Checklist

• Doc strings have been added in the Numpy docstring format.
  Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.
• All linting and tests pass.

[orionarcher]

LGTM, just request a couple small things. Tagging @abhijeetgangan for a look too.

[orionarcher]

The stress tensor should have shape n_batches, 3, 3, would you convert this to follow that convention?

[stefanbringuier]

Sure, no problem.

[orionarcher]

Could you add a test to check that if I run unbatched(A) and unbatched(B) I get the same results as batched(A, B)?

[abhijeetgangan]

• Currently, not all models return per-atom potential energies and stresses. It is out of our control if the original implementation of the model does not return those properties. The kinetic part is easy to add.
• One needs to be careful about the volume normalization. This is not present in the LAMMPS per-atom stresses so you might need to multiple by the volume.
• You can have a test case with LJ as it can return the per-atom properties.

Does this definition work for potentials that use angles? Since we are using auto-grad to get the stresses it will include all the terms beyond pair-wise interactions. LAMMPS has a warning about that.

[stefanbringuier]

@abhijeetgangan thanks for the comments and look-outs!

• Currently, not all models return per-atom potential energies and stresses. It is out of our control if the original implementation of the model does not return those properties. The kinetic part is easy to add.

I thought this might be an issue for general model use, any thoughts on how to handle?

• One needs to be careful about the volume normalization. This is not present in the LAMMPS per-atom stresses so you might need to multiple by the volume.

Yeah, the formalism here does not include the domain volume normalization intentionally because I was intending to use it HFACF and for scenarios where the heatflux is calculated for a subset of atoms rather than the whole domain. But that was probably not a suitable reason to exclude it. Can add option to return domain volume normalized J.

• You can have a test case with LJ as it can return the per-atom properties.

Working on a example script that was this.

Does this definition work for potentials that use angles? Since we are using auto-grad to get the stresses it will include all the terms beyond pair-wise interactions. LAMMPS has a warning about that.

Hmm, no haven't tested and need to review. The tests I've implemented just check formalism. Thanks!

[abhijeetgangan]

• @stefanbringuier MACE returns per-atom potential energies, and it should be possible to get per-atom stresses. However, that means modifying the Shift-Scale MACE or doing the auto-grad in the forward in torch-sim. I am unsure what is the best way to deal with it is, but I will think about it over the weekend. We do want to use this with ML potentials, after all.

• Regarding points 2-3 - Do you have a good reference value from literature with LJ to get the Heat flux / Kappa? That would make a robust test as well. I will also search for it or maybe a similar LAMMPS example.

• We can safely use it for pairwise interaction potentials, I think. Need to be careful if the formalism itself is limiting for models like mattersim which use angles.

Thank you for the PR. This goes well with the correlation PR for Kappa calculations. I was planning to add it myself but wasn't sure because of the limitations I mentioned above.

[stefanbringuier]

@abhijeetgangan

Did some searching and found a paper by Admal & Tadmor (2015) where they derive heat flux expressions without requiring per-atom energies, using the energy balance law and force–position coupling. But the key issue is that it requires access to partial forces $f_{ij}$, which MLIPs don't typically expose to my knowledge.

Then I found Carbogno et al. (2017), where they adopt a simplified version using only the virial term: $J_v = \frac{1}{V} \sum_i \sigma_i \cdot v_i$, with $\sigma_i \sim r_i \otimes f_i$. Their argument is that for solids the convective term is negligible, so this form gives valid Green–Kubo $\kappa$. This avoids needing per-atom energies or per-atom stresses.

So my thinking is: when per-atom stresses aren’t provided, we extend calc_heat_flux to optionally accept per-atom positions and forces and construct position-force contractions internally. Then if energies=None, we just return the virial part with a warning that the convective term is assumed negligible.

Let me know what you think? Will proceed if you're on board with that.

[abhijeetgangan]

@stefanbringuier all good points. I have tested here some of them before.

I also merged the LJ to include per-atom quantities in a batch so you can proceed with a example script for both cases where you have per atom quantities vs positions-forces contractions following the script above.

I think there is a better way to make this approach more general but for now this should be good enough for a test case. This reference has a detailed derivation of the virial. I would also include the references you mentioned.

[janosh]

what's the state of this PR? happy to review if it's ready

[stefanbringuier]

what's the state of this PR? happy to review if it's ready

It's still not quite ready as I need to review @abhijeetgangan tests and also work on an example/test using the LJ per-atom quantities vs. the contraction calculation. Will prioritize to bring the PR to a close (been focused on the NEB feature).

[janosh]

ofc, no rush. happy to leave this open. just checking it's not held up by us