Skip to content

Add a heat flux wrapper#144

Open
GardevoirX wants to merge 33 commits intometatensor:mainfrom
GardevoirX:heat-flux-wrapper
Open

Add a heat flux wrapper#144
GardevoirX wants to merge 33 commits intometatensor:mainfrom
GardevoirX:heat-flux-wrapper

Conversation

@GardevoirX
Copy link
Contributor

@GardevoirX GardevoirX commented Jan 26, 2026 •
edited
Loading

Closes #104

This is still in a very early stage and subjected to change.

TODO:

  • Interface for getting the heat flux from the model
  • Checks for the applicability of unfolding the system
  • Tests
  • Documentations
  • Remove Replace the dependency on vesin.metatomic with vesin.torch and rebuild neighbor list from the un-unfolded system (Hard to build the neighboring information between the replica atoms, e.g. atoms in different image boxes. And vesin.metatomic is not supposed to be used in metatomic.)

Contributor (creator of pull-request) checklist

  • Tests updated (for new features and bugfixes)?
  • Documentation updated (for new features)?
  • Issue referenced (for PRs that solve an issue)?

Reviewer checklist

  • CHANGELOG updated with public API or any other important changes?

Luthaf
Luthaf reviewed Jan 28, 2026
View reviewed changes
@GardevoirX GardevoirX force-pushed the heat-flux-wrapper branch 2 times, most recently from 35c184e to d684615 Compare February 3, 2026 15:17
This was referenced Feb 12, 2026
Merged
Merged
@GardevoirX GardevoirX marked this pull request as ready for review February 24, 2026 22:00
@GardevoirX
Copy link
Contributor Author

GardevoirX commented Feb 24, 2026 •
edited
Loading

I think this is now ready for reviewing. Though the current results with pet models are not good enough, soap-bpnn still gives some promising results. Plus, there's no evident bug in the current implementation. Perhaps a thorough checking for the code can reveal underlying problems.

In terms of the implementation, thie current state hasn't reached an ideal state. The calculation of the term1 of the heat flux can be implemented in a JVP style, but it is still in VJP and calculated through a for-loop. This might be improved after metatensor/metatensor#1063 is merged, but previous tests suggested that more blocks are in the path.

The second thing is how we do "unfolding" in _unfold_system. There could be two ways:

  1. The current way: check if particles are close to the box boundary, and generate the replicas if true.
  2. Do a NL calculation to get the atom pairs closer than the interaction range of the model, and then only unfolding the atoms in the pairs with non-zero shifting vector.

Tests showed that the second way can generate fewer atoms and thus can save runtime from the model inference, while an extra NL calculation also costs time. My intuition is that which one is faster depends on the system size.

Also maybe we need a tutorial about how to use this wrapper.

Luthaf
Luthaf reviewed Feb 27, 2026
View reviewed changes
@GardevoirX GardevoirX force-pushed the heat-flux-wrapper branch 2 times, most recently from f44cf5c to b1b81da Compare March 9, 2026 15:27
@GardevoirX GardevoirX requested a review from Luthaf March 9, 2026 22:56
Luthaf
Luthaf reviewed Mar 12, 2026
View reviewed changes
python/metatomic_torch/tests/heat_flux.py Outdated
Comment on lines +366 to +368
def test_heat_flux_wrapper_rejects_non_eV_energy(model_in_kcal_per_mol):
with pytest.raises(ValueError, match="energy outputs in eV"):
HeatFluxWrapper(model_in_kcal_per_mol)
Copy link
Member

@Luthaf Luthaf Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I don't love this but we need #173 first

@GardevoirX GardevoirX mentioned this pull request Mar 12, 2026
Open
4 tasks
GardevoirX and others added 17 commits March 20, 2026 10:49
@GardevoirX
Replace self._model = model with self._model = model.module and r…
b870b71 
…emove some tests
@GardevoirX
Minor fix
191c934
@GardevoirX
Remove skin
cae3d36
@GardevoirX
Add heat flux as a standard output, also add momentum and velocity as…
acf39a4 
… outputs of i-pi

Fix CI

Try again

Again

Add an example on how to use the wrapper

Update documentation

Try to fix docs again

Hopefully the last fix
@GardevoirX
Support energy variant
f002d1d
@GardevoirX
Add a test for wrapper asking for the same quantity as the inner mode…
e179760 
…l but with different unit
@GardevoirX
Check the length unit of inner model
f2ea03e
@GardevoirX @Luthaf
Update python/metatomic_torch/metatomic/torch/model.py
3c7bf95 
Co-authored-by: Guillaume Fraux <luthaf@luthaf.fr>
@GardevoirX
Trace models that ask for the same quantity but in different units
3c4ba5f
@GardevoirX
Only keep smoke tests
09094ad
@Luthaf @GardevoirX
Minor cleanup
2867c00
@Luthaf @GardevoirX
Remove torch.jit.export
83b6105
@GardevoirX
Renaming and add wrap method to help user create heat flux model ea…
915e051 
…sily
@GardevoirX
Needs grads
94d2a22
@GardevoirX
Replace Angstrom with angstrom
5e61727
@GardevoirX
Fix CI
eb0b5cc
@GardevoirX
Reuse NL calculators
7257c39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Luthaf Luthaf Luthaf left review comments

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Implement a wrapper for heat flux calculation

2 participants

@GardevoirX @Luthaf