Update to metatrain v2026.4 and fix uncertainty under rotational averaging - #161
Open
ppegolo wants to merge 13 commits into
Open
Update to metatrain v2026.4 and fix uncertainty under rotational averaging#161ppegolo wants to merge 13 commits into
ppegolo wants to merge 13 commits into
Conversation
metatrain v2026.4 requires metatomic-torch >=0.1.17, which deprecates two ModelOutput arguments this package still passed: - `quantity` is deprecated outright, so it is dropped; - `per_atom` is replaced by `sample_kind`, which takes "atom"/"system" rather than a boolean. The `per_atom` flag stays in this package's own public signatures (`get_energy_uncertainty`, `get_energy_ensemble`, `_calculate_dos`); only the translation into a `ModelOutput` changes. `metatomic-ase` is pinned to the >=0.1.3 release that goes with metatomic-torch 0.1.17.
`rotational_average_order` replaces `self.calculator` with a `SymmetrizedCalculator`, which has neither `_calculate_uncertainty` nor `_energy_uq_key` nor `run_model`: every uncertainty method raised an `AttributeError` as soon as rotational averaging was enabled. No test combined the two, so this went unnoticed. Uncertainty outputs are not rotationally averaged, so they are requested from the wrapped calculator, reached through a `_base_calculator` property that unwraps `SymmetrizedCalculator` when there is one.
Register pet-omol-{s,m,l} as available models, and give them their own
metadata: the authors and references of the model paper
(doi:10.1088/2632-2153/ae6417), plus the OMol25 dataset paper. The
dataset reference is listed under `model` because `references` only
accepts the `implementation`, `architecture` and `model` keys.
Register pet-mols-s as an available model, and as a UQ-capable one since the released checkpoints are wrapped in an LLPR uncertainty model. It is not a universal potential, so it gets its own description on top of its own authors and reference (arXiv:2603.06236).
ppegolo
marked this pull request as ready for review
August 13, 2026 12:56
The test suite is parametrized over UPET_AVAILABLE_MODELS, so registering PET-OMOL and PET-MOLS enrolled them in tests written around materials models, and 11 of them failed. The non-conservative regime used to require both non-conservative forces and stresses. Models trained on non-periodic data (PET-OMOL) only predict the forces, so they are now run with the forces alone, leaving the stress to backpropagation. PET-MOLS has no non-conservative outputs at all and is listed as such. The remaining failures came from the tests themselves: they assumed every non-MAD model shares one set of authors and a universal-potential description, they hardcoded PET-SPICE as the only model without non-conservative outputs, and they evaluated every model on stretched diamond, which is far enough out of PET-MOLS' domain that its uncertainty exceeds the threshold above which the calculator warns.
Delete UPET_NO_NC_STRESS_MODELS: it was never used, and the calculator already reads the same fact from the model capabilities. Inline the PET-OMOL and PET-MOLS metadata in their branches, like the other models, instead of hoisting them to module constants and looking the dataset name up in a one-entry dict. The metadata test used to reproduce the branches of get_upet_metadata, so it could only fail if one of the two copies was edited on its own. Test one model per family against literal values instead. Also document the new models, and report the missing non-conservative forces rather than "forces and stresses", since the stress alone is now tolerated.
It followed the list of conservative-only checkpoints, and so read as an explanation of the PET-MOLS entries in it.
abmazitov
reviewed
Aug 14, 2026
Comment on lines
+126
to
+127
| PET-OMol models only predict non-conservative forces, so their stress | ||
| is computed by backpropagation even when this is set. |
Collaborator
There was a problem hiding this comment.
This looks a bit confusing, how can a model for molecules have no conservative forces, but at the same time have conservative stresses?
Contributor
Author
There was a problem hiding this comment.
You can still use it in the bulk if you want, and in that case it's fine to compute (conservative) stress. But this comment is confusing, so I'll get rid of it.
abmazitov
reviewed
Aug 14, 2026
Comment on lines
+179
to
+181
| if nc_stress_key not in model_outputs: | ||
| # the stress is left to backpropagation | ||
| nc_regime = "forces" |
Collaborator
There was a problem hiding this comment.
Why do we need to fall back to this option? Our models typically either have or don't have the nc forces and stresses. Besides, I'm not a particular fan of having such implicit silent swaps of a model behavior
Contributor
Author
There was a problem hiding this comment.
It's not silent because we just checked nc forces are available in the line above. It isn't a fallback, it is just a flag saying such a model only supports nc forces (like OMol models). Happy to discuss a better design if you dislike this one
Removed outdated information about PET-OMol models in the docstring.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The
metatrainpin no longer resolved against a released version. Bumping it tov2026.4 pulls in metatomic-torch 0.1.17, which deprecates two
ModelOutputarguments we passed:
quantityis dropped, andper_atombecomessample_kind.Our own
per_atomarguments stay as they are.It also fixes the uncertainty methods when rotational averaging is on. They were
called on the
SymmetrizedCalculator, which does not have the attributes theyneed, so they raised an
AttributeErroras soon asrotational_average_orderwas set. They now go to the calculator it wraps, since uncertainties are not
rotationally averaged. Nothing tested the two together, hence the silence.