Skip to content

Commit fdc99e1

Browse files
maxbortoneclaudelwalew
authored
feat: add NVE energy-conservation benchmark (#146)
* feat: add NVE energy-conservation benchmark Port the NVE energy-conservation benchmark from the internal repo. It runs short microcanonical (NVE, velocity-Verlet, no thermostat) MD trajectories for four representative systems (vacuum molecule, bulk water, two solvated peptides) and scores how well the model conserves total mechanical energy E(t) = PE(t) + KE(t). The metric is the fitted total-energy drift over the run divided by the kinetic-energy standard deviation, mapped through the standard soft-threshold scorer. - Add the benchmark module, registered and gated per-system on elements. - Add the Streamlit UI page (per-system drift curves + linear fits) and wire it into the GUI. - Mark it beta (BENCHMARKS_TO_SKIP_FOR_PUBLIC_LEADERBOARD), like the NEB benchmark. - Add unit tests and checked-in test structures, plus the generic UI-page test. - Add user-facing and API-reference docs under the General category. The NVE velocity-Verlet integrator and SimulationState.potential_energy are only on mlip's `develop` branch, not yet in a PyPI release, so pin mlip to git `develop` via [tool.uv.sources] until a release ships them. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore: address review comments on NVE benchmark - Drop the temporary mlip git pin now that MDIntegrator.NVE_VELOCITY_VERLET and SimulationState.potential_energy have shipped in mlip 0.2.2; bump the dependency floor to >=0.2.2 (both core and cuda extra) so resolution can't fall back to a release missing these features. - Display the NVE energy-conservation benchmark on the public leaderboard. - Document the ENERGY_DRIFT_RATIO_THRESHOLD = 1.0 rationale and why it is more lenient than the ~0.01 used for classical force fields. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: prevent NVE Time (ps) axis title from being clipped The total-energy-drift chart used a fixed height with container-width rendering, which clipped the x-axis title. Size the chart to fit within its height (including padding) so the "Time (ps)" title is fully shown. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: add CHANGELOG entry for the NVE energy-conservation benchmark Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> # Conflicts: # CHANGELOG.md * fix: axis cutoff by legend --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: lwalew <l.walewski@instadeep.com>
1 parent 78ec2a1 commit fdc99e1

17 files changed

Lines changed: 6630 additions & 2 deletions

File tree

CHANGELOG.md

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
11
# Changelog
22

3+
## Release 0.1.5
4+
5+
- Add the `nve_energy_conservation` benchmark, which runs short NVE (constant-energy)
6+
MD simulations across a set of gas-phase and solvated systems and scores how well a
7+
model conserves the total energy, quantified by the drift of the total energy over
8+
the trajectory.
9+
310
## Release 0.1.4
411

512
- Add an Apache-2.0 `LICENSE` file and declare the license in `pyproject.toml`.
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
.. _nve_energy_conservation_api:
2+
3+
NVE energy conservation
4+
=======================
5+
6+
.. module:: mlipaudit.benchmarks.nve_energy_conservation.nve_energy_conservation
7+
8+
.. autoclass:: NVEEnergyConservationBenchmark
9+
10+
.. automethod:: __init__
11+
12+
.. automethod:: run_model
13+
14+
.. automethod:: analyze
15+
16+
.. autoclass:: NVEEnergyConservationResult
17+
18+
.. autoclass:: NVEStructureResult
19+
20+
.. autoclass:: NVEEnergyConservationModelOutput

docs/source/api_reference/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,4 @@ Benchmark implementations
4040
biomolecules/sampling
4141
general/stability
4242
general/scaling
43+
general/nve_energy_conservation

docs/source/benchmarks/general/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,4 @@ applicable across molecular systems.
1111

1212
Stability <stability>
1313
Scaling <scaling>
14+
NVE energy conservation <nve_energy_conservation>
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
.. _nve_energy_conservation:
2+
3+
NVE energy conservation
4+
=======================
5+
6+
Purpose
7+
-------
8+
9+
To assess how well a machine-learned interatomic potential (**MLIP**) conserves the total
10+
mechanical energy during microcanonical (**NVE**) molecular dynamics (**MD**). Energy
11+
conservation is a fundamental physical requirement: under **NVE** dynamics, with no
12+
thermostat exchanging energy with the environment, the total energy
13+
:math:`E(t) = \mathrm{PE}(t) + \mathrm{KE}(t)` should stay constant. A systematic drift
14+
indicates unphysical forces or an inconsistency between the model's energy and its
15+
gradient.
16+
17+
Description
18+
-----------
19+
20+
For each system in the dataset, the benchmark runs a short **NVE** (velocity-Verlet, no
21+
thermostat) **MD** simulation with the **MLIP**, with velocities initialized from a
22+
Maxwell-Boltzmann distribution at **300 K**, leveraging
23+
`jax-md <https://github.com/google/jax-md>`_ as integrated via the
24+
`mlip <https://github.com/instadeepai/mlip>`_ library. The potential and kinetic energies
25+
are recorded at regular snapshots, and the total-energy drift relative to the first frame,
26+
:math:`\Delta E(t) = E(t) - E(0)`, is fitted with a linear model.
27+
28+
The **energy-conservation metric** is the magnitude of the fitted drift over the whole
29+
trajectory, normalized by the kinetic-energy fluctuation scale:
30+
31+
.. math::
32+
33+
r = \frac{\lvert \text{slope} \rvert \cdot T}{\sigma_{\mathrm{KE}}}
34+
35+
where :math:`T` is the trajectory duration and :math:`\sigma_{\mathrm{KE}}` is the
36+
standard deviation of the kinetic energy along the trajectory. This dimensionless ratio is
37+
mapped to a score in :math:`[0, 1]` via the standard soft threshold: a ratio at or below
38+
**1.0** scores **1.0**, and larger ratios decay exponentially.
39+
40+
The threshold of **1.0** is set at the point where the systematic energy drift accumulated
41+
over the whole trajectory reaches the same magnitude as the natural kinetic-energy
42+
fluctuations: a ratio :math:`r \le 1` keeps the drift within the thermal noise floor, while
43+
:math:`r > 1` means the drift dominates over those fluctuations. This is a deliberately
44+
lenient threshold. Analogous energy-conservation tests for classical force fields typically
45+
use a far tighter ratio, on the order of **0.01**; the more permissive value here reflects
46+
a fundamental difference between the two. A classical force field is an analytic expression
47+
whose potential-energy surface is smooth by construction, whereas an **MLIP** is a highly
48+
parameterized model whose potential-energy surface is learned from finite training data.
49+
While the model is trained to reproduce the training data, it is not guaranteed to be
50+
smooth in between training points, as the interpolation between them can exhibit
51+
high-frequency roughness. This roughness manifests as steep, sudden gradients, which the
52+
MD integrator can struggle with. The resulting numerical error can accumulate and drive
53+
larger integration drift compared to a classical force field. The threshold of **1.0** is
54+
therefore chosen to be lenient enough to accommodate this behavior, while still being
55+
strict enough to flag models that are clearly unphysical.
56+
57+
Dataset
58+
-------
59+
60+
The dataset comprises four representative systems spanning vacuum, bulk liquid and solvated
61+
regimes:
62+
63+
- Small molecule (HCNO) in vacuum
64+
- Bulk water (500 molecules)
65+
- Solvated peptide (Oxytocin)
66+
- Solvated peptide with counter-ions (Neurotensin)
67+
68+
Systems whose elements the model cannot handle are skipped individually rather than
69+
skipping the whole benchmark.
70+
71+
Interpretation
72+
--------------
73+
74+
A score of **1.0** indicates excellent energy conservation, i.e. the total-energy drift is
75+
small relative to the natural kinetic-energy fluctuations. A score approaching **0.0**
76+
indicates a strongly drifting, non-conservative trajectory.

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ requires-python = ">=3.11"
2727
dependencies = [
2828
"huggingface-hub>=0.33.4",
2929
"mdtraj>=1.10.3",
30-
"mlip>=0.2.0,<0.3.0",
30+
"mlip>=0.2.2,<0.3.0",
3131
"scikit-learn>=1.7.0",
3232
"streamlit>=1.46.1",
3333
"vl-convert-python>=1.8.0",
@@ -45,7 +45,7 @@ build-backend = "hatchling.build"
4545

4646
[project.optional-dependencies]
4747
cuda = [
48-
"mlip[cuda13]>=0.2.0,<0.3.0",
48+
"mlip[cuda13]>=0.2.2,<0.3.0",
4949
]
5050

5151
[dependency-groups]

src/mlipaudit/benchmarks/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,12 @@
4444
NEBResult,
4545
NudgedElasticBandBenchmark,
4646
)
47+
from mlipaudit.benchmarks.nve_energy_conservation.nve_energy_conservation import (
48+
NVEEnergyConservationBenchmark,
49+
NVEEnergyConservationModelOutput,
50+
NVEEnergyConservationResult,
51+
NVEStructureResult,
52+
)
4753
from mlipaudit.benchmarks.reactivity.reactivity import (
4854
ReactivityBenchmark,
4955
ReactivityModelOutput,
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2025 InstaDeep Ltd
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

0 commit comments

Comments
 (0)