feat: add NVE energy-conservation benchmark

docs/source/api_reference/general/nve_energy_conservation.rst

@@ -0,0 +1,20 @@
+.. _nve_energy_conservation_api:
+
+NVE energy conservation
+=======================
+
+.. module:: mlipaudit.benchmarks.nve_energy_conservation.nve_energy_conservation
+
+.. autoclass:: NVEEnergyConservationBenchmark
+
+    .. automethod:: __init__
+
+    .. automethod:: run_model
+
+    .. automethod:: analyze
+
+.. autoclass:: NVEEnergyConservationResult
+
+.. autoclass:: NVEStructureResult
+
+.. autoclass:: NVEEnergyConservationModelOutput

docs/source/api_reference/index.rst

@@ -40,3 +40,4 @@ Benchmark implementations
     biomolecules/sampling
     general/stability
     general/scaling
+    general/nve_energy_conservation

docs/source/benchmarks/general/index.rst

@@ -11,3 +11,4 @@ applicable across molecular systems.
 
     Stability <stability>
     Scaling <scaling>
+    NVE energy conservation <nve_energy_conservation>

docs/source/benchmarks/general/nve_energy_conservation.rst

@@ -0,0 +1,59 @@
+.. _nve_energy_conservation:
+
+NVE energy conservation
+=======================
+
+Purpose
+-------
+
+To assess how well a machine-learned interatomic potential (**MLIP**) conserves the total
+mechanical energy during microcanonical (**NVE**) molecular dynamics (**MD**). Energy
+conservation is a fundamental physical requirement: under **NVE** dynamics, with no
+thermostat exchanging energy with the environment, the total energy
+:math:`E(t) = \mathrm{PE}(t) + \mathrm{KE}(t)` should stay constant. A systematic drift
+indicates unphysical forces or an inconsistency between the model's energy and its
+gradient.
+
+Description
+-----------
+
+For each system in the dataset, the benchmark runs a short **NVE** (velocity-Verlet, no
+thermostat) **MD** simulation with the **MLIP**, with velocities initialized from a
+Maxwell-Boltzmann distribution at **300 K**, leveraging
+`jax-md <https://github.com/google/jax-md>`_ as integrated via the
+`mlip <https://github.com/instadeepai/mlip>`_ library. The potential and kinetic energies
+are recorded at regular snapshots, and the total-energy drift relative to the first frame,
+:math:`\Delta E(t) = E(t) - E(0)`, is fitted with a linear model.
+
+The **energy-conservation metric** is the magnitude of the fitted drift over the whole
+trajectory, normalized by the kinetic-energy fluctuation scale:
+
+.. math::
+
+   r = \frac{\lvert \text{slope} \rvert \cdot T}{\sigma_{\mathrm{KE}}}
+
+where :math:`T` is the trajectory duration and :math:`\sigma_{\mathrm{KE}}` is the
+standard deviation of the kinetic energy along the trajectory. This dimensionless ratio is
+mapped to a score in :math:`[0, 1]` via the standard soft threshold: a ratio at or below
+**1.0** scores **1.0**, and larger ratios decay exponentially.
+
+Dataset
+-------
+
+The dataset comprises four representative systems spanning vacuum, bulk liquid and solvated
+regimes:
+
+   - Small molecule (HCNO) in vacuum
+   - Bulk water (500 molecules)
+   - Solvated peptide (Oxytocin)
+   - Solvated peptide with counter-ions (Neurotensin)
+
+Systems whose elements the model cannot handle are skipped individually rather than
+skipping the whole benchmark.
+
+Interpretation
+--------------
+
+A score of **1.0** indicates excellent energy conservation, i.e. the total-energy drift is
+small relative to the natural kinetic-energy fluctuations. A score approaching **0.0**
+indicates a strongly drifting, non-conservative trajectory.

pyproject.toml

@@ -31,6 +31,13 @@ cuda = [
     "mlip[cuda13]>=0.2.0,<0.3.0",
 ]
 
+# Temporary: the NVE energy-conservation benchmark needs the NVE velocity-Verlet
+# integrator (MDIntegrator.NVE_VELOCITY_VERLET) and SimulationState.potential_energy,
+# which are on mlip's `develop` branch but not yet in a published PyPI release. Pin to
+# git `develop` until a release ships these, then drop this source to use the PyPI build.
+[tool.uv.sources]
+mlip = { git = "https://github.com/instadeepai/mlip-jax.git", branch = "develop" }
+
 [dependency-groups]
 dev = [
     "notebook>=7.4.4",

src/mlipaudit/benchmarks/__init__.py

@@ -44,6 +44,12 @@
     NEBResult,
     NudgedElasticBandBenchmark,
 )
+from mlipaudit.benchmarks.nve_energy_conservation.nve_energy_conservation import (
+    NVEEnergyConservationBenchmark,
+    NVEEnergyConservationModelOutput,
+    NVEEnergyConservationResult,
+    NVEStructureResult,
+)
 from mlipaudit.benchmarks.reactivity.reactivity import (
     ReactivityBenchmark,
     ReactivityModelOutput,
@@ -97,7 +103,10 @@
 BENCHMARKS_WITHOUT_SCORES = [ScalingBenchmark]
 
 # Some benchmarks are still in beta and are not displayed in the public leaderboard
-BENCHMARKS_TO_SKIP_FOR_PUBLIC_LEADERBOARD = [NudgedElasticBandBenchmark]
+BENCHMARKS_TO_SKIP_FOR_PUBLIC_LEADERBOARD = [
+    NudgedElasticBandBenchmark,
+    NVEEnergyConservationBenchmark,
+]
 
 
 def _setup_benchmark_categories() -> dict[str, list[type[Benchmark]]]:

src/mlipaudit/benchmarks/nve_energy_conservation/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 InstaDeep Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.