feat: split molecular liquids into RDF + density benchmarks

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Changelog
 
+## Release 0.1.5
+
+- Split the Molecular Liquids benchmarks into four separately scored benchmarks:
+  `water_radial_distribution`, `water_density`, `solvent_radial_distribution`, and
+  `solvent_density`, all grouped under the `"Molecular Liquids"` category
+- Add a `data_name` attribute (and `data_dir` property) to `Benchmark` so that
+  several benchmarks can share the same input data directory and HuggingFace archive.
+- Add `Water density` and `Solvent density` UI pages showing the equilibrium-density
+  summary statistics and a per-frame density time series against the experimental
+  reference.
+
 ## Release 0.1.3
 
 - Populate `atoms.info["charge"]` and `atoms.info["spin"]` on every benchmark's

docs/source/api_reference/index.rst

@@ -32,10 +32,12 @@ Benchmark implementations
     small_molecules/ring_planarity
     small_molecules/reference_geometry_stability
     small_molecules/bond_length_distribution
-    small_molecules/radial_distribution
-    small_molecules/solvent_radial_distribution
     small_molecules/reactivity
     small_molecules/nudged_elastic_band
+    molecular_liquids/water_radial_distribution
+    molecular_liquids/solvent_radial_distribution
+    molecular_liquids/water_density
+    molecular_liquids/solvent_density
     biomolecules/folding_stability
     biomolecules/sampling
     general/stability

docs/source/api_reference/molecular_liquids/solvent_density.rst

@@ -0,0 +1,20 @@
+.. _solvent_density_api:
+
+Solvent Density
+===============
+
+.. module:: mlipaudit.benchmarks.solvent_density.solvent_density
+
+.. autoclass:: SolventDensityBenchmark
+
+    .. automethod:: __init__
+
+    .. automethod:: run_model
+
+    .. automethod:: analyze
+
+.. autoclass:: SolventDensityResult
+
+.. autoclass:: SolventDensityStructureResult
+
+.. autoclass:: SolventDensityModelOutput

docs/source/api_reference/molecular_liquids/water_density.rst

@@ -0,0 +1,18 @@
+.. _water_density_api:
+
+Water Density
+=============
+
+.. module:: mlipaudit.benchmarks.water_density.water_density
+
+.. autoclass:: WaterDensityBenchmark
+
+    .. automethod:: __init__
+
+    .. automethod:: run_model
+
+    .. automethod:: analyze
+
+.. autoclass:: WaterDensityResult
+
+.. autoclass:: WaterDensityModelOutput

docs/source/benchmarks/molecular_liquids/density.rst

@@ -0,0 +1,59 @@
+.. _density:
+
+Density
+=======
+
+Purpose
+-------
+
+This benchmark assesses the ability of machine-learned interatomic potentials (**MLIP**) to
+reproduce the **equilibrium density** of molecular liquids. Density is a fundamental
+thermodynamic property: a model that predicts accurate local structure (see the
+:ref:`radial_distribution` benchmark) may still get the density badly wrong if the simulation
+box expands or collapses. Reproducing the correct density is therefore a complementary and
+necessary check on the physical realism of a liquid-phase simulation.
+
+Description
+-----------
+
+The benchmark runs the same **MD** simulation as the :ref:`radial_distribution` benchmark: an
+**NPT** simulation using the **MLIP** model for **500,000 steps**, leveraging the
+`jax-md <https://github.com/google/jax-md>`_ engine from the
+`mlip <https://github.com/instadeepai/mlip>`_ library. Water is run at **295.15 K** and **1 atm**,
+while all other solvents are run at **293.15 K** and **1 atm**. Because the :ref:`radial_distribution` and density
+benchmarks of a system share their input systems and simulation output, the simulation is only
+run once when both benchmarks are run together.
+
+The density of each frame is computed from the (fluctuating) simulation cell volume:
+
+.. math::
+
+   \rho = \frac{N_\text{mol} \, M}{N_A \, V}
+
+where :math:`N_\text{mol}` is the number of molecules in the box, :math:`M` is the molecular
+weight, :math:`N_A` is Avogadro's number and :math:`V` is the cell volume. The **equilibrium
+density** is taken as the average density over the final four fifths of the trajectory (the
+first fifth is discarded as equilibration).
+
+Dataset
+-------
+
+The benchmark uses the same equilibrated input boxes as the :ref:`radial_distribution` benchmark
+(a 500-molecule TIP3P water box, and methanol / acetonitrile / CCl4 boxes built with the GAFF
+force field in OpenMM).
+
+Reference densities are the experimental values at the simulation conditions: water
+:math:`0.9978\ \text{g/cm}^3`, CCl4 :math:`1.594\ \text{g/cm}^3`, methanol
+:math:`0.791\ \text{g/cm}^3` and acetonitrile :math:`0.786\ \text{g/cm}^3`.
+
+Interpretation
+--------------
+
+Performance is quantified by the **density deviation**, the absolute difference between the
+equilibrium density and the experimental reference. The deviation should be **as low as
+possible**. The score is derived from the *relative* deviation (deviation divided by the
+reference density) so that it is comparable across liquids of very different densities: a
+relative deviation within roughly **2%** scores close to 1, decaying gently beyond that. (The
+ideal per-solvent target is ultimately set by the isothermal compressibility of the liquid.) A
+large deviation typically indicates that the box has expanded or collapsed during the
+simulation, and the density time series can be inspected on the results page to diagnose this.

docs/source/benchmarks/molecular_liquids/index.rst

@@ -4,10 +4,11 @@ Molecular Liquids
 =================
 
 Molecular Liquids benchmarks are focused on the properties and dynamics of molecular liquids,
-including as initial benchmark its radial distribution function.
+including their radial distribution function and equilibrium density.
 
 
 .. toctree::
     :maxdepth: 1
 
     Radial distribution function <radial_distribution>
+    Density <density>

docs/source/benchmarks/molecular_liquids/radial_distribution.rst

@@ -21,9 +21,8 @@ leveraging the `jax-md <https://github.com/google/jax-md>`_ engine from the
 `mlip <https://github.com/instadeepai/mlip>`_ library. Water is run at **295.15 K** and **1 atm**,
 while all other solvents are run at **293.15 K** and **1 atm**. The starting configuration is already
 equilibrated. For every specific atom pair (e.g., **oxygen-oxygen** in water) the radial distribution
-function (**RDF** or **g(r)**) is calculated from the simulation.
-
-# TODO: Add details of density computation + details RE experimental density to sections below
+function (**RDF** or **g(r)**) is calculated from the simulation. The equilibrium
+density of the same simulation is assessed separately in the :ref:`density` benchmark.
 
 .. figure:: img/rdf.png
     :figwidth: 35%
@@ -55,9 +54,9 @@ experimental reference data. Performance is quantified using the following metri
 
 Dataset
 -------
-For the water radial distribution benchmark we set up a cubic box of 500 water molecules using OpenMM and the TIP3P water model.
+For the :ref:`water radial distribution benchmark <water_radial_distribution_api>` we set up a cubic box of 500 water molecules using OpenMM and the TIP3P water model.
 We equilibrated the box in the NPT ensemble at standard conditions and extracted the final snapshot as input for the benchmark.
-For the solvent radial distribution benchmark, we initialized the solvent boxes (methanol, acetonitrile, CCl4) by stacking randomly rotated molecules
+For the :ref:`solvent radial distribution benchmark <solvent_radial_distribution_api>`, we initialized the solvent boxes (methanol, acetonitrile, CCl4) by stacking randomly rotated molecules
 to yield a cubic box with a target side-length of 28 Å at the experimental density. We equilibrated the box in the NPT ensemble using the GAFF force field and OpenMM.
 
 We use the experimental water RDF profile of Skinner et al.\ [#f1]_ as reference data. For other solvents (methanol\ [#f2]_, acetonitrile\ [#f3]_, CCl4\ [#f4]_), we use the

src/mlipaudit/benchmark.py

@@ -93,6 +93,10 @@ class Benchmark(ABC):
             If present, a user or the CLI can make use of this information to reuse
             cached model outputs from another benchmark carrying the same ID instead of
             rerunning simulations or inference.
+        data_name: An optional name of the input data directory (and HuggingFace
+            archive) to use instead of `name`. This lets several benchmarks share the
+            same input data without duplicating it. Defaults to None, in which case
+            `name` is used.
     """
 
     name: str = ""
@@ -105,6 +109,8 @@ class Benchmark(ABC):
 
     reusable_output_id: tuple[str, ...] | None = None
 
+    data_name: str | None = None
+
     def __init__(
         self,
         force_field: ForceField | ASECalculator,
@@ -134,9 +140,9 @@ def __init__(
                 required elements.
             ValueError: If force field type is not compatible.
         """
-        self.run_mode = run_mode
-        if not isinstance(self.run_mode, RunMode):
-            self.run_mode = RunMode(run_mode)
+        self.run_mode: RunMode = (
+            run_mode if isinstance(run_mode, RunMode) else RunMode(run_mode)
+        )
 
         self.force_field = force_field
 
@@ -225,17 +231,27 @@ def check_can_run_model(cls, force_field: ForceField) -> bool:
 
         return True
 
+    @property
+    def data_dir(self) -> Path:
+        """The local directory holding this benchmark's input data.
+
+        Uses `data_name` when set, otherwise `name`, so that benchmarks can share
+        input data.
+        """
+        return self.data_input_dir / (self.data_name or self.name)
+
     def _download_data(self) -> None:
         """Download the data from the data input directory if not already exists."""
-        already_exists = (self.data_input_dir / self.name).exists()
+        data_name = self.data_name or self.name
+        already_exists = (self.data_input_dir / data_name).exists()
         if not already_exists:
             hf_hub_download(
                 repo_id="InstaDeepAI/MLIPAudit-data",
-                filename=f"{self.name}.zip",
+                filename=f"{data_name}.zip",
                 local_dir=self.data_input_dir,
                 repo_type="dataset",
             )
-            with zipfile.ZipFile(self.data_input_dir / f"{self.name}.zip", "r") as z:
+            with zipfile.ZipFile(self.data_input_dir / f"{data_name}.zip", "r") as z:
                 z.extractall(self.data_input_dir)
 
     @abstractmethod

src/mlipaudit/benchmarks/__init__.py

@@ -69,6 +69,12 @@
     ScalingModelOutput,
     ScalingResult,
 )
+from mlipaudit.benchmarks.solvent_density.solvent_density import (
+    SolventDensityBenchmark,
+    SolventDensityModelOutput,
+    SolventDensityResult,
+    SolventDensityStructureResult,
+)
 from mlipaudit.benchmarks.solvent_radial_distribution.solvent_radial_distribution import (  # noqa: E501
     SolventRadialDistributionBenchmark,
     SolventRadialDistributionModelOutput,
@@ -85,6 +91,11 @@
     TautomersModelOutput,
     TautomersResult,
 )
+from mlipaudit.benchmarks.water_density.water_density import (
+    WaterDensityBenchmark,
+    WaterDensityModelOutput,
+    WaterDensityResult,
+)
 from mlipaudit.benchmarks.water_radial_distribution.water_radial_distribution import (
     WaterRadialDistributionBenchmark,
     WaterRadialDistributionModelOutput,