Conversation
ppegolo
force-pushed
the
forces_uq
branch
2 times, most recently
from
July 9, 2026 13:51
1d9139c to
cb5e5c9
Compare
ppegolo
marked this pull request as draft
July 28, 2026 14:53
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.
ppegolo
marked this pull request as ready for review
August 13, 2026 07:09
ppegolo
changed the title
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).
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.
Checkpoints now store their own name, description, authors and references, so a new model does not need an entry here: PET-OMol and PET-MOLS get theirs from the checkpoints published for them, and only the models released before that are still described here. The metadata used to be passed to the export, which is also where the model gives its own: an uncertainty model merges the metadata of the model it wraps last, so the empty metadata of an older wrapped checkpoint would overwrite the one passed here, and PET-MAD v1.5.0 was exported without a name or a description at all. Fill in the metadata of the model that has none instead, which leaves nothing to overwrite.
Whether a model predicts non-conservative forces or an uncertainty on its energy is written in its capabilities, and was read from there already: the two lists of model versions kept next to them were only used to build an error message and to tell the tests what to expect, and had to be edited by hand for every model published. Expose the two capabilities as properties of the calculator instead, so that the answer comes from the model and callers can ask for it before running.
A model offered here but missing from the hub only fails when someone tries to download it, which is late. The other direction is left alone: a model can be published before it is offered here.
Uncertainty estimates were previously limited to the energy. This adds them for forces and stresses, via two routes that the calculator exposes through `get_forces_ensemble` / `get_forces_uncertainty`, the stress equivalents, and `get_forces_and_stress_ensemble`: - "conservative": the energy ensemble is differentiated with respect to positions and cell, so each member's forces are the gradient of that member's energy and are conservative by construction. The per-member mean over atoms is removed to restore translational invariance. Forces and stresses share a single backward pass when asked for together. - "direct": models carrying a separately-trained vectorial force head expose `mtt::aux::non_conservative_forces_ensemble` and the matching `_uncertainty` output, sampled in that head's last-layer weight space. This is much cheaper than differentiating the ensemble, but the forces are not conservative. The two are genuinely different estimates, not two implementations of one, so the method is selectable. It defaults to whichever matches the regime the calculator was built in, and in conservative mode the direct ensemble is re-centered on the conservative forces so that only the spread differs. Stress has no "direct" counterpart: models do not carry an ensemble for `non_conservative_stress`.
The list of models with uncertainty estimates is gone, so the errors of the force and stress ensembles now say what the documentation says, and the calculator is asked whether the model provides an uncertainty at all.
ppegolo
changed the base branch from
update-metatrain-2026.4
to
derive-model-info
August 13, 2026 15:26
The calculator was two thirds uncertainty estimates, most of it plumbing: requesting outputs from the model and reshaping what comes back. That part now lives in `_uncertainty`, taking what it needs as arguments, and the calculator is left with the questions only it can answer: whether the model provides these outputs, and which atoms to run on.
…aging (#161) * Update to metatrain v2026.4 and migrate deprecated metatomic APIs 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. * Fix the uncertainty methods under rotational averaging `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. * Add the PET-OMOL models 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. * Add the PET-MOLS model 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). * Make the test suite fit the new molecular models 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. * Clean up the metadata and non-conservative changes 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. * Drop the note about the non-conservative stress of PET-OMol It followed the list of conservative-only checkpoints, and so read as an explanation of the PET-MOLS entries in it. * Update docstring to remove PET-OMol note Removed outdated information about PET-OMol models in the docstring. * Simplified the logic for non-conservative evaluation * Fixed a little error and tests * Fixed tests * One more attempt to fix tests * One more attempt to fix tests --------- Co-authored-by: Arslan Mazitov <arslan.mazitov@phystech.edu>
* Take the model metadata from the checkpoint Checkpoints now store their own name, description, authors and references, so a new model does not need an entry here: PET-OMol and PET-MOLS get theirs from the checkpoints published for them, and only the models released before that are still described here. The metadata used to be passed to the export, which is also where the model gives its own: an uncertainty model merges the metadata of the model it wraps last, so the empty metadata of an older wrapped checkpoint would overwrite the one passed here, and PET-MAD v1.5.0 was exported without a name or a description at all. Fill in the metadata of the model that has none instead, which leaves nothing to overwrite. * Check the available models against the hub A model offered here but missing from the hub only fails when someone tries to download it, which is late. The other direction is left alone: a model can be published before it is offered here. * Derive non-conservative support from the model and rework its evaluation Whether a model predicts non-conservative forces or an uncertainty on its energy is written in its capabilities, and was read from there already: the two lists of model versions kept next to them were only used to build an error message and to tell the tests what to expect, and had to be edited by hand for every model published. Expose the two capabilities as properties of the calculator instead, so that the answer comes from the model and callers can ask for it before running, and collect the available non-conservative quantities once so the requested regime is checked against them. Co-authored-by: ppegolo <paolo.pegolo@epfl.ch> * Linting fix * Adapted for yanked metatomic-ase versions * Constrained nvalchemi-toolkit version * Added back the nvalchemi-toolkit-ops version restriction * Updated the docs --------- Co-authored-by: Arslan Mazitov <arslan.mazitov@phystech.edu>
Resolve the overlap with the non-conservative rework that landed on main: - The NC setup, its error message and the tests asserting on it come from main, which superseded the earlier `_supports_non_conservative`/`nc_regime` version carried here; the silent forces-only fallback is gone with it, so an unsupported request raises instead of degrading. - `supports_non_conservative` now reads main's `_available_nc_quantities`. - The uncertainty machinery, its move to `_uncertainty.py` and the force and stress uncertainty tests are kept from this branch. - pyproject keeps the higher `metatomic-ase` floor from here and main's `nvalchemi-toolkit-ops` constraint.
`get_forces_ensemble` and `get_forces_uncertainty` take a `non_conservative` flag again, defaulting to the regime the calculator was built with. The direct ensemble is re-centered on the conservative forces when the calculator is conservative, which is what makes "direct force uncertainty" the spread of the direct head around the forces actually in use. Asking a non-conservative calculator for a conservative ensemble is an error, as the result would not match the forces it returns. Two fixes come with it: - the uncertainty and ensemble outputs of any target other than a plain `energy` are named with an `mtt::aux::` prefix by metatrain, not only when a variant is selected. Without this the non-conservative UQ outputs are unreachable on a checkpoint with no variants. - `get_forces_uncertainty` is the spread of the corresponding ensemble, so it refers to the same forces the calculator returns. The model's own direct uncertainty is computed before the net force is removed, and that mode carries ~80% of the head's variance, so it overestimates the error on the projected forces by ~2x. It stays as a fallback for models carrying a direct uncertainty but no direct ensemble.
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.
Uncertainty estimates were limited to the energy. This adds them for forces and
stresses, through
get_forces_uncertainty,get_stress_uncertainty, their_ensemblecounterparts, andget_forces_and_stress_ensemble, which computesboth in one backward pass.
Forces come either from differentiating the energy ensemble, which keeps them
conservative, or from a direct force head for the models that have one, which is
much cheaper but not conservative. The two are different estimates rather than
two ways of getting one, so the method is selectable, and defaults to whichever
matches how the calculator was built. Stress only has the first: no model
carries an ensemble for
non_conservative_stress.No checkpoint on HuggingFace has a direct force head yet, so that route only
works with locally trained models.
One of the points to discuss: whether to actually default to NC-force uncertainty, when
available, as conservative ones are very expensive