From 4777de22131595b2356d2f208a5f06ec61eaa1a4 Mon Sep 17 00:00:00 2001 From: lwalew Date: Thu, 2 Jul 2026 16:32:21 +0200 Subject: [PATCH 01/12] feat: split molecular liquids RDF and density into separate benchmarks Split the water and solvent Molecular Liquids benchmarks into four scored benchmarks: water/solvent x RDF/density. Each is its own Benchmark subclass, score, leaderboard column and UI page, grouped under the "Molecular Liquids" category. - The RDF and density benchmark of each system group share a single NPT simulation via `reusable_output_id` and identical `ModelOutput` field sets, so the simulation is only run once when both are run together. - Density is now scored on its own (relative deviation from the reference density via `compute_metric_score`), rather than only shown in the RDF UI. - Shared simulation config, data loaders and density helpers moved to `utils/molecular_liquids.py`. - New `Benchmark.data_name` lets the density benchmarks reuse the RDF input data without duplicating HuggingFace uploads. - New density UI pages plot the density time series against the reference. - Docs: added a density benchmark page and moved the RDF api-reference docs under `api_reference/molecular_liquids/`. Note: the density score threshold (DENSITY_RELATIVE_DEVIATION_THRESHOLD) is a placeholder pending sign-off from the science team. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/source/api_reference/index.rst | 6 +- .../molecular_liquids/solvent_density.rst | 20 ++ .../solvent_radial_distribution.rst | 0 .../molecular_liquids/water_density.rst | 18 ++ .../water_radial_distribution.rst} | 0 .../benchmarks/molecular_liquids/density.rst | 57 ++++ .../benchmarks/molecular_liquids/index.rst | 3 +- .../molecular_liquids/radial_distribution.rst | 5 +- src/mlipaudit/benchmark.py | 22 +- src/mlipaudit/benchmarks/__init__.py | 11 + .../solvent_density/solvent_density.py | 220 +++++++++++++ .../solvent_radial_distribution.py | 139 +------- .../benchmarks/water_density/water_density.py | 157 +++++++++ .../water_radial_distribution.py | 107 +------ src/mlipaudit/ui/__init__.py | 2 + src/mlipaudit/ui/solvent_density.py | 182 +++++++++++ .../ui/solvent_radial_distribution.py | 5 - src/mlipaudit/ui/water_density.py | 171 ++++++++++ src/mlipaudit/ui/water_radial_distribution.py | 3 - src/mlipaudit/utils/molecular_liquids.py | 299 +++++++++++++++++- tests/solvent_density/test_solvent_density.py | 136 ++++++++ .../test_solvent_radial_distribution.py | 4 - tests/test_ui_pages.py | 24 ++ tests/water_density/test_water_density.py | 133 ++++++++ .../test_water_radial_distribution.py | 4 - 25 files changed, 1485 insertions(+), 243 deletions(-) create mode 100644 docs/source/api_reference/molecular_liquids/solvent_density.rst rename docs/source/api_reference/{small_molecules => molecular_liquids}/solvent_radial_distribution.rst (100%) create mode 100644 docs/source/api_reference/molecular_liquids/water_density.rst rename docs/source/api_reference/{small_molecules/radial_distribution.rst => molecular_liquids/water_radial_distribution.rst} (100%) create mode 100644 docs/source/benchmarks/molecular_liquids/density.rst create mode 100644 src/mlipaudit/benchmarks/solvent_density/solvent_density.py create mode 100644 src/mlipaudit/benchmarks/water_density/water_density.py create mode 100644 src/mlipaudit/ui/solvent_density.py create mode 100644 src/mlipaudit/ui/water_density.py create mode 100644 tests/solvent_density/test_solvent_density.py create mode 100644 tests/water_density/test_water_density.py diff --git a/docs/source/api_reference/index.rst b/docs/source/api_reference/index.rst index 34655795..4bf3f218 100644 --- a/docs/source/api_reference/index.rst +++ b/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 diff --git a/docs/source/api_reference/molecular_liquids/solvent_density.rst b/docs/source/api_reference/molecular_liquids/solvent_density.rst new file mode 100644 index 00000000..20fb97eb --- /dev/null +++ b/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 diff --git a/docs/source/api_reference/small_molecules/solvent_radial_distribution.rst b/docs/source/api_reference/molecular_liquids/solvent_radial_distribution.rst similarity index 100% rename from docs/source/api_reference/small_molecules/solvent_radial_distribution.rst rename to docs/source/api_reference/molecular_liquids/solvent_radial_distribution.rst diff --git a/docs/source/api_reference/molecular_liquids/water_density.rst b/docs/source/api_reference/molecular_liquids/water_density.rst new file mode 100644 index 00000000..4b70a99b --- /dev/null +++ b/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 diff --git a/docs/source/api_reference/small_molecules/radial_distribution.rst b/docs/source/api_reference/molecular_liquids/water_radial_distribution.rst similarity index 100% rename from docs/source/api_reference/small_molecules/radial_distribution.rst rename to docs/source/api_reference/molecular_liquids/water_radial_distribution.rst diff --git a/docs/source/benchmarks/molecular_liquids/density.rst b/docs/source/benchmarks/molecular_liquids/density.rst new file mode 100644 index 00000000..836138da --- /dev/null +++ b/docs/source/benchmarks/molecular_liquids/density.rst @@ -0,0 +1,57 @@ +.. _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 `_ engine from the +`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 RDF 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 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. diff --git a/docs/source/benchmarks/molecular_liquids/index.rst b/docs/source/benchmarks/molecular_liquids/index.rst index 3b809524..02c25991 100644 --- a/docs/source/benchmarks/molecular_liquids/index.rst +++ b/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 + Density diff --git a/docs/source/benchmarks/molecular_liquids/radial_distribution.rst b/docs/source/benchmarks/molecular_liquids/radial_distribution.rst index 7598fe4f..f7119f3f 100644 --- a/docs/source/benchmarks/molecular_liquids/radial_distribution.rst +++ b/docs/source/benchmarks/molecular_liquids/radial_distribution.rst @@ -21,9 +21,8 @@ leveraging the `jax-md `_ engine from the `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% diff --git a/src/mlipaudit/benchmark.py b/src/mlipaudit/benchmark.py index 502a55d6..c94a355b 100644 --- a/src/mlipaudit/benchmark.py +++ b/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, @@ -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 (e.g. an RDF and a density benchmark running the same system). + """ + 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 diff --git a/src/mlipaudit/benchmarks/__init__.py b/src/mlipaudit/benchmarks/__init__.py index 5e899c38..bc22e461 100644 --- a/src/mlipaudit/benchmarks/__init__.py +++ b/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, diff --git a/src/mlipaudit/benchmarks/solvent_density/solvent_density.py b/src/mlipaudit/benchmarks/solvent_density/solvent_density.py new file mode 100644 index 00000000..16ae23fa --- /dev/null +++ b/src/mlipaudit/benchmarks/solvent_density/solvent_density.py @@ -0,0 +1,220 @@ +# 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. +import logging +import statistics + +import numpy as np +from mlip.simulation import SimulationState +from pydantic import BaseModel, ConfigDict, NonNegativeFloat + +from mlipaudit.benchmark import ( + Benchmark, + BenchmarkResult, + ModelOutput, +) +from mlipaudit.scoring import ALPHA, compute_metric_score +from mlipaudit.utils.molecular_liquids import ( + DENSITY_RELATIVE_DEVIATION_THRESHOLD, + SOLVENT_DATA_NAME, + SOLVENT_MOLECULE_CONFIG, + SOLVENT_REFERENCE_DENSITIES, + SOLVENT_REUSABLE_OUTPUT_ID, + average_equilibrated_density, + compute_densities, + run_solvent_npt_simulations, +) +from mlipaudit.utils.stability import is_simulation_stable + +logger = logging.getLogger("mlipaudit") + + +class SolventDensityModelOutput(ModelOutput): + """Model output containing the final simulation states for each structure. + + Shares its field signature with `SolventRadialDistributionModelOutput` so that the + NPT simulations can be reused between the two benchmarks (see `reusable_output_id`). + + Attributes: + structure_names: The names of the structures. + simulation_states: `SimulationState` or `None` object for each structure in + the same order as the structure names. `None` if the simulation failed. + """ + + structure_names: list[str] + simulation_states: list[SimulationState | None] + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class SolventDensityStructureResult(BaseModel): + """Stores the density result for a single structure. + + Attributes: + structure_name: The structure name. + densities: List of per-frame densities in g/cm3. + average_density: Average density over the final four fifths of the frames. + density_deviation: Absolute deviation of the average density from the reference. + reference_density: The reference density in g/cm3. + failed: Whether the simulation was successful. If unsuccessful, the other + attributes will be not be set. + score: The score for the molecule. + """ + + structure_name: str + densities: list[float] | None = None + average_density: float | None = None + density_deviation: NonNegativeFloat | None = None + reference_density: float | None = None + + failed: bool = False + score: float = 0.0 + + +class SolventDensityResult(BenchmarkResult): + """Result object for the solvent density benchmark. + + Attributes: + structure_names: The names of the structures. + structures: List of per structure results. + avg_density_deviation: The average density deviation across all structures. + failed: Whether all the simulations failed and no analysis could be + performed. Defaults to False. + score: The final score for the benchmark between 0 and 1. + """ + + structure_names: list[str] + structures: list[SolventDensityStructureResult] + avg_density_deviation: NonNegativeFloat | None = None + + +class SolventDensityBenchmark(Benchmark): + """Benchmark for the equilibrium density of molecular solvents. + + Runs the same NPT simulations as `SolventRadialDistributionBenchmark` (reusing + their output when both are run together) and scores how closely the equilibrium + density of each solvent matches its reference. + + Attributes: + name: The unique benchmark name. The name is `solvent_density`. + category: The benchmark's category, `"Molecular Liquids"`. + data_name: The input-data directory shared with the solvent RDF benchmark. + result_class: `SolventDensityResult`. + model_output_class: `SolventDensityModelOutput`. + required_elements: The set of atomic element types present in the input files. + reusable_output_id: Shared with `SolventRadialDistributionBenchmark` so the + solvent NPT simulations are only run once when both benchmarks are run + together. + """ + + name = "solvent_density" + category = "Molecular Liquids" + data_name = SOLVENT_DATA_NAME + result_class = SolventDensityResult + model_output_class = SolventDensityModelOutput + + required_elements = {"N", "H", "O", "C", "Cl"} + + reusable_output_id = SOLVENT_REUSABLE_OUTPUT_ID + + def run_model(self) -> None: + """Run an MD simulation for each structure using the NPT ensemble. + + The MD simulation is performed using the JAX MD engine and starts from + the reference structure. The NPT integrator uses Langevin dynamics with + a Monte Carlo barostat. + """ + structure_names, simulation_states = run_solvent_npt_simulations( + self.force_field, self.data_dir, self.run_mode + ) + self.model_output = SolventDensityModelOutput( + structure_names=structure_names, simulation_states=simulation_states + ) + + def analyze(self) -> SolventDensityResult: + """Compute the equilibrium density of each solvent and its deviation. + + Returns: + A `SolventDensityResult` object. + + Raises: + RuntimeError: If called before `run_model()`. + """ + if self.model_output is None: + raise RuntimeError("Must call run_model() first.") + + structure_results = [] + num_succeeded = 0 + + for system_name, simulation_state in zip( + self.model_output.structure_names, self.model_output.simulation_states + ): + if simulation_state is None or not is_simulation_stable(simulation_state): + structure_results.append( + SolventDensityStructureResult( + structure_name=system_name, + failed=True, + score=0.0, + ) + ) + continue + + num_succeeded += 1 + mol_config = SOLVENT_MOLECULE_CONFIG[system_name] + densities = compute_densities( + simulation_state, + mol_config["molecule_weight"], + int(mol_config["atoms_per_molecule"]), + ) + average_density = average_equilibrated_density(densities) + reference_density = SOLVENT_REFERENCE_DENSITIES[system_name] + density_deviation = abs(average_density - reference_density) + + relative_deviation = density_deviation / reference_density + score = compute_metric_score( + np.array([relative_deviation]), + DENSITY_RELATIVE_DEVIATION_THRESHOLD, + ALPHA, + ).item() + + structure_results.append( + SolventDensityStructureResult( + structure_name=system_name, + densities=densities.tolist(), + average_density=average_density, + density_deviation=density_deviation, + reference_density=reference_density, + score=score, + ) + ) + + if num_succeeded == 0: + return SolventDensityResult( + structure_names=self.model_output.structure_names, + structures=structure_results, + failed=True, + score=0.0, + ) + + return SolventDensityResult( + structure_names=self.model_output.structure_names, + structures=structure_results, + avg_density_deviation=statistics.mean( + structure.density_deviation + for structure in structure_results + if structure.density_deviation is not None + ), + score=statistics.mean( + r.score if r.score is not None else 0.0 for r in structure_results + ), + ) diff --git a/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py b/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py index 175c5f26..5572264f 100644 --- a/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py +++ b/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py @@ -17,54 +17,27 @@ import mdtraj as md import numpy as np -from ase import Atoms, units -from ase.io import read as ase_read +from ase import units from mlip.simulation import SimulationState -from mlip.simulation.enums import MDIntegrator from pydantic import BaseModel, ConfigDict, NonNegativeFloat from mlipaudit.benchmark import ( - DEFAULT_CHARGE, - DEFAULT_SPIN, Benchmark, BenchmarkResult, ModelOutput, ) -from mlipaudit.run_mode import RunMode from mlipaudit.scoring import ALPHA -from mlipaudit.utils import ( - create_mdtraj_trajectory_from_simulation_state, - run_simulation, +from mlipaudit.utils import create_mdtraj_trajectory_from_simulation_state +from mlipaudit.utils.molecular_liquids import ( + SOLVENT_REUSABLE_OUTPUT_ID, + get_solvent_pdb_file_name, + get_solvent_system_names, + run_solvent_npt_simulations, ) -from mlipaudit.utils.molecular_liquids import compute_densities from mlipaudit.utils.stability import is_simulation_stable logger = logging.getLogger("mlipaudit") -SIMULATION_CONFIG = { - "num_steps": 500_000, - "snapshot_interval": 500, - "num_episodes": 1000, - "temperature_kelvin": 293.15, - "pressure_bar": 1.01325, -} - -SIMULATION_CONFIG_DEV = { - "num_steps": 5, - "snapshot_interval": 1, - "num_episodes": 1, - "temperature_kelvin": 293.15, - "pressure_bar": 1.01325, -} -SIMULATION_CONFIG_FAST = { - "num_steps": 250_000, - "snapshot_interval": 250, - "num_episodes": 1000, - "temperature_kelvin": 293.15, - "pressure_bar": 1.01325, -} -NUM_DEV_SYSTEMS = 1 - SYSTEM_ATOM_OF_INTEREST = { "CCl4": "C", "methanol": "O", @@ -73,11 +46,6 @@ MIN_RADII, MAX_RADII = 0.0, 20.0 # In Angstrom -MOLECULE_CONFIG = { - "CCl4": {"molecule_weight": 153.823, "atoms_per_molecule": 5}, - "methanol": {"molecule_weight": 32.042, "atoms_per_molecule": 6}, - "acetonitrile": {"molecule_weight": 41.053, "atoms_per_molecule": 6}, -} REFERENCE_MAXIMA = { "CCl4": {"type": "C-C", "distance": 5.9, "range": (0.0, 20.0)}, "acetonitrile": {"type": "N-N", "distance": 4.0, "range": (3.5, 4.5)}, @@ -89,12 +57,6 @@ "methanol": (0.0, 20.0), } -REFERENCE_DENSITIES = { - "CCl4": 1.594, - "acetonitrile": 0.786, - "methanol": 0.791, -} - class SolventRadialDistributionModelOutput(ModelOutput): """Model output containing the final simulation states for @@ -118,9 +80,6 @@ class SolventRadialDistributionStructureResult(BaseModel): Attributes: structure_name: The structure name. - densities: List of densities in g/cm3. - average_density: Average density over the final 4 fifths of the frames. - density_deviation: Deviation of the average density from the reference. radii: The radii values in Angstrom. rdf: The radial distribution function values at the radii. first_solvent_peak: The first solvent peak, i.e. the radius at which the @@ -132,9 +91,6 @@ class SolventRadialDistributionStructureResult(BaseModel): """ structure_name: str - densities: list[float] | None = None - average_density: float | None = None - density_deviation: NonNegativeFloat | None = None radii: list[float] | None = None rdf: list[float] | None = None first_solvent_peak: float | None = None @@ -150,7 +106,6 @@ class SolventRadialDistributionResult(BenchmarkResult): Attributes: structure_names: The names of the structures. structures: List of per structure results. - avg_density_deviation: The average density deviation across all structures. avg_peak_deviation: The average peak deviation across all structures. failed: Whether all the simulations failed and no analysis could be performed. Defaults to False. @@ -160,7 +115,6 @@ class SolventRadialDistributionResult(BenchmarkResult): structure_names: list[str] structures: list[SolventRadialDistributionStructureResult] - avg_density_deviation: NonNegativeFloat | None = None avg_peak_deviation: NonNegativeFloat | None = None @@ -185,6 +139,8 @@ class SolventRadialDistributionBenchmark(Benchmark): if there are some atomic element types that the model cannot handle. If False, the benchmark must have its own custom logic to handle missing atomic element types. For this benchmark, the attribute is set to True. + reusable_output_id: Shared with `SolventDensityBenchmark` so that the solvent + NPT simulations are only run once when both benchmarks are run together. """ name = "solvent_radial_distribution" @@ -194,6 +150,8 @@ class SolventRadialDistributionBenchmark(Benchmark): required_elements = {"N", "H", "O", "C", "Cl"} + reusable_output_id = SOLVENT_REUSABLE_OUTPUT_ID + def run_model(self) -> None: """Run an MD simulation for each structure using the NPT ensemble. @@ -201,28 +159,11 @@ def run_model(self) -> None: the reference structure. The NPT integrator uses Langevin dynamics with a Monte Carlo barostat. """ - if self.run_mode == RunMode.DEV: - md_kwargs = SIMULATION_CONFIG_DEV - elif self.run_mode == RunMode.FAST: - md_kwargs = SIMULATION_CONFIG_FAST - else: - md_kwargs = SIMULATION_CONFIG - - simulation_states = [] - for system_name in self._system_names: - logger.info("Running MD for %s radial distribution function.", system_name) - - simulation_state = run_simulation( - atoms=self._load_system(system_name), - force_field=self.force_field, - md_integrator=MDIntegrator.NPT_MC_LANGEVIN, - molecule_indices=self._load_molecule_indices(system_name), - **md_kwargs, - ) - simulation_states.append(simulation_state) - + structure_names, simulation_states = run_solvent_npt_simulations( + self.force_field, self.data_dir, self.run_mode + ) self.model_output = SolventRadialDistributionModelOutput( - structure_names=self._system_names, simulation_states=simulation_states + structure_names=structure_names, simulation_states=simulation_states ) def analyze(self) -> SolventRadialDistributionResult: @@ -254,21 +195,10 @@ def analyze(self) -> SolventRadialDistributionResult: continue num_succeeded += 1 - mol_config = MOLECULE_CONFIG[system_name] - densities = compute_densities( - simulation_state, - mol_config["molecule_weight"], - int(mol_config["atoms_per_molecule"]), - ) - n_frames_equilibration = len(densities) // 5 - average_density = np.mean(densities[n_frames_equilibration:]) - density_deviation = abs(average_density - REFERENCE_DENSITIES[system_name]) traj = create_mdtraj_trajectory_from_simulation_state( simulation_state=simulation_state, - topology_path=self.data_input_dir - / self.name - / self._get_pdb_file_name(system_name), + topology_path=self.data_dir / get_solvent_pdb_file_name(system_name), ) pair_indices = traj.top.select( f"symbol == {SYSTEM_ATOM_OF_INTEREST[system_name]}" @@ -313,9 +243,6 @@ def analyze(self) -> SolventRadialDistributionResult: structure_result = SolventRadialDistributionStructureResult( structure_name=system_name, - densities=densities, - average_density=average_density, - density_deviation=density_deviation, radii=radii.tolist(), rdf=rdf, first_solvent_peak=first_solvent_peak, @@ -336,11 +263,6 @@ def analyze(self) -> SolventRadialDistributionResult: return SolventRadialDistributionResult( structure_names=self.model_output.structure_names, structures=structure_results, - avg_density_deviation=statistics.mean( - structure.density_deviation - for structure in structure_results - if structure.density_deviation is not None - ), avg_peak_deviation=statistics.mean( structure.peak_deviation for structure in structure_results @@ -353,31 +275,4 @@ def analyze(self) -> SolventRadialDistributionResult: @property def _system_names(self) -> list[str]: - if self.run_mode == RunMode.STANDARD: - return list(SYSTEM_ATOM_OF_INTEREST.keys()) - - # reduced number of cases for DEV and FAST run mode - return list(SYSTEM_ATOM_OF_INTEREST.keys())[:NUM_DEV_SYSTEMS] - - def _load_system(self, system_name) -> Atoms: - atoms = ase_read( - self.data_input_dir / self.name / self._get_pdb_file_name(system_name) - ) - atoms.info["charge"] = DEFAULT_CHARGE - atoms.info["spin"] = DEFAULT_SPIN - return atoms - - def _load_molecule_indices(self, system_name) -> np.ndarray: - molecule_indices_filename = self._get_molecule_indices_file_name(system_name) - molecule_indices = np.load( - self.data_input_dir / self.name / molecule_indices_filename - ) - return molecule_indices - - @staticmethod - def _get_pdb_file_name(system_name: str) -> str: - return f"{system_name}_eq.pdb" - - @staticmethod - def _get_molecule_indices_file_name(system_name: str) -> str: - return f"{system_name}_molecule_indices.npy" + return get_solvent_system_names(self.run_mode) diff --git a/src/mlipaudit/benchmarks/water_density/water_density.py b/src/mlipaudit/benchmarks/water_density/water_density.py new file mode 100644 index 00000000..cff2e988 --- /dev/null +++ b/src/mlipaudit/benchmarks/water_density/water_density.py @@ -0,0 +1,157 @@ +# 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. +import logging + +import numpy as np +from mlip.simulation import SimulationState +from pydantic import ConfigDict, NonNegativeFloat + +from mlipaudit.benchmark import ( + Benchmark, + BenchmarkResult, + ModelOutput, +) +from mlipaudit.scoring import ALPHA, compute_metric_score +from mlipaudit.utils.molecular_liquids import ( + DENSITY_RELATIVE_DEVIATION_THRESHOLD, + WATER_ATOMS_PER_MOLECULE, + WATER_DATA_NAME, + WATER_MOLECULE_WEIGHT, + WATER_REFERENCE_DENSITY, + WATER_REUSABLE_OUTPUT_ID, + average_equilibrated_density, + compute_densities, + run_water_npt_simulation, +) +from mlipaudit.utils.stability import is_simulation_stable + +logger = logging.getLogger("mlipaudit") + + +class WaterDensityModelOutput(ModelOutput): + """Model output containing the final simulation state of the water box. + + Shares its field signature with `WaterRadialDistributionModelOutput` so that the + NPT simulation can be reused between the two benchmarks (see `reusable_output_id`). + + Attributes: + simulation_state: The final simulation state of the water + box simulation. None if the simulation failed. + failed: Whether the simulation failed. Defaults to False. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + simulation_state: SimulationState | None = None + failed: bool = False + + +class WaterDensityResult(BenchmarkResult): + """Result object for the water density benchmark. + + Attributes: + densities: List of per-frame densities in g/cm3. + average_density: Average density over the final four fifths of the frames. + density_deviation: Absolute deviation of the average density from the + experimental reference. + reference_density: The experimental reference density in g/cm3. + failed: Whether the simulation failed and no analysis could be + performed. Defaults to False. + score: The final score for the benchmark between 0 and 1. + """ + + densities: list[float] | None = None + average_density: float | None = None + density_deviation: NonNegativeFloat | None = None + reference_density: float = WATER_REFERENCE_DENSITY + + +class WaterDensityBenchmark(Benchmark): + """Benchmark for the equilibrium density of liquid water. + + Runs the same water box NPT simulation as `WaterRadialDistributionBenchmark` + (reusing its output when both are run together) and scores how closely the + equilibrium density matches the experimental reference. + + Attributes: + name: The unique benchmark name. The name is `water_density`. + category: The benchmark's category, `"Molecular Liquids"`. + data_name: The input-data directory shared with the water RDF benchmark. + result_class: `WaterDensityResult`. + model_output_class: `WaterDensityModelOutput`. + required_elements: The set of atomic element types present in the input files. + reusable_output_id: Shared with `WaterRadialDistributionBenchmark` so the water + box NPT simulation is only run once when both benchmarks are run together. + """ + + name = "water_density" + category = "Molecular Liquids" + data_name = WATER_DATA_NAME + result_class = WaterDensityResult + model_output_class = WaterDensityModelOutput + + required_elements = {"H", "O"} + + reusable_output_id = WATER_REUSABLE_OUTPUT_ID + + def run_model(self) -> None: + """Run an MD simulation for the water box system using the NPT ensemble. + + The MD simulation is performed using the JAX MD engine and starts from + the reference structure. The NPT integrator uses Langevin dynamics with + a Monte Carlo barostat. + """ + simulation_state = run_water_npt_simulation( + self.force_field, self.data_dir, self.run_mode + ) + self.model_output = WaterDensityModelOutput( + simulation_state=simulation_state, failed=simulation_state is None + ) + + def analyze(self) -> WaterDensityResult: + """Compute the equilibrium density and how much it deviates from the reference. + + Returns: + A `WaterDensityResult` object. + + Raises: + RuntimeError: If called before `run_model()`. + """ + if self.model_output is None: + raise RuntimeError("Must call run_model() first.") + + simulation_state = self.model_output.simulation_state + + if self.model_output.failed or not is_simulation_stable(simulation_state): + return WaterDensityResult(failed=True, score=0.0) + + densities = compute_densities( + simulation_state, WATER_MOLECULE_WEIGHT, WATER_ATOMS_PER_MOLECULE + ) + average_density = average_equilibrated_density(densities) + density_deviation = abs(average_density - WATER_REFERENCE_DENSITY) + + relative_deviation = density_deviation / WATER_REFERENCE_DENSITY + score = compute_metric_score( + np.array([relative_deviation]), + DENSITY_RELATIVE_DEVIATION_THRESHOLD, + ALPHA, + ).item() + + return WaterDensityResult( + densities=densities.tolist(), + average_density=average_density, + density_deviation=density_deviation, + score=score, + ) diff --git a/src/mlipaudit/benchmarks/water_radial_distribution/water_radial_distribution.py b/src/mlipaudit/benchmarks/water_radial_distribution/water_radial_distribution.py index a35dd237..fc6537a2 100644 --- a/src/mlipaudit/benchmarks/water_radial_distribution/water_radial_distribution.py +++ b/src/mlipaudit/benchmarks/water_radial_distribution/water_radial_distribution.py @@ -11,31 +11,27 @@ # 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. -import functools import logging import math -from typing import Any import mdtraj as md import numpy as np -from ase import Atoms, units -from ase.io import read as ase_read +from ase import units from mlip.simulation import SimulationState -from mlip.simulation.enums import MDIntegrator from pydantic import ConfigDict, NonNegativeFloat from sklearn.metrics import mean_absolute_error, root_mean_squared_error from mlipaudit.benchmark import ( - DEFAULT_CHARGE, - DEFAULT_SPIN, Benchmark, BenchmarkResult, ModelOutput, ) -from mlipaudit.run_mode import RunMode from mlipaudit.scoring import ALPHA, compute_metric_score -from mlipaudit.utils import run_simulation -from mlipaudit.utils.molecular_liquids import compute_densities +from mlipaudit.utils.molecular_liquids import ( + WATER_REUSABLE_OUTPUT_ID, + WATERBOX_N500, + run_water_npt_simulation, +) from mlipaudit.utils.stability import is_simulation_stable from mlipaudit.utils.trajectory_helpers import ( create_mdtraj_trajectory_from_simulation_state, @@ -43,43 +39,13 @@ logger = logging.getLogger("mlipaudit") -SIMULATION_CONFIG = { - "num_steps": 500_000, - "snapshot_interval": 500, - "num_episodes": 1000, - "temperature_kelvin": 295.15, - "pressure_bar": 1.01325, -} - -SIMULATION_CONFIG_FAST = { - "num_steps": 250_000, - "snapshot_interval": 250, - "num_episodes": 1000, - "temperature_kelvin": 295.15, - "pressure_bar": 1.01325, -} - -SIMULATION_CONFIG_DEV = { - "num_steps": 5, - "snapshot_interval": 1, - "num_episodes": 1, - "temperature_kelvin": 295.15, - "pressure_bar": 1.01325, -} - -WATERBOX_N500 = "water_box_n500_eq.pdb" -MOLECULE_INDICES_PATH = "water_box_n500_molecule_indices.npy" REFERENCE_DATA = "experimental_reference.npz" -MOLECULE_WEIGHT = 18.01528 # g/mol -ATOMS_PER_MOLECULE = 3 REFERENCE_PEAK_DISTANCE = 2.80 # A RMSE_SCORE_THRESHOLD = 0.1 SOLVENT_PEAK_RANGE = (2.8, 3.0) RADII_RANGE = (2.5, 10.0) -REFERENCE_DENSITY = 0.997773 # g/cm3 - class WaterRadialDistributionModelOutput(ModelOutput): """Model output containing the final simulation state of the water box. @@ -100,9 +66,6 @@ class WaterRadialDistributionResult(BenchmarkResult): """Result object for the water radial distribution benchmark. Attributes: - densities: List of densities in g/cm3. - average_density: Average density over the final 4 fifths of the frames. - density_deviation: Deviation of the average density from the reference. radii: The radii values in Angstrom. rdf: The radial distribution function values at the radii. mae: The MAE of the radial distribution function values. @@ -118,9 +81,6 @@ class WaterRadialDistributionResult(BenchmarkResult): score: The final score for the benchmark between 0 and 1. """ - densities: list[float] | None = None - average_density: float | None = None - density_deviation: NonNegativeFloat | None = None radii: list[float] | None = None rdf: list[float] | None = None mae: float | None = None @@ -151,6 +111,8 @@ class WaterRadialDistributionBenchmark(Benchmark): if there are some atomic element types that the model cannot handle. If False, the benchmark must have its own custom logic to handle missing atomic element types. For this benchmark, the attribute is set to True. + reusable_output_id: Shared with `WaterDensityBenchmark` so that the water box + NPT simulation is only run once when both benchmarks are run together. """ name = "water_radial_distribution" @@ -160,6 +122,8 @@ class WaterRadialDistributionBenchmark(Benchmark): required_elements = {"H", "O"} + reusable_output_id = WATER_REUSABLE_OUTPUT_ID + def run_model(self) -> None: """Run an MD simulation for the water box system using the NPT ensemble. @@ -167,16 +131,9 @@ def run_model(self) -> None: the reference structure. The NPT integrator uses Langevin dynamics with a Monte Carlo barostat. """ - logger.info("Running MD for water radial distribution function.") - - simulation_state = run_simulation( - atoms=self._water_box_n500, - force_field=self.force_field, - md_integrator=MDIntegrator.NPT_MC_LANGEVIN, - molecule_indices=self._molecule_indices, - **self._md_kwargs, + simulation_state = run_water_npt_simulation( + self.force_field, self.data_dir, self.run_mode ) - self.model_output = WaterRadialDistributionModelOutput( simulation_state=simulation_state, failed=simulation_state is None ) @@ -198,16 +155,9 @@ def analyze(self) -> WaterRadialDistributionResult: if self.model_output.failed or not is_simulation_stable(simulation_state): return WaterRadialDistributionResult(failed=True, score=0.0) - densities = compute_densities( - simulation_state, MOLECULE_WEIGHT, ATOMS_PER_MOLECULE - ) - n_frames_equilibration = len(densities) // 5 - average_density = np.mean(densities[n_frames_equilibration:]) - density_deviation = abs(average_density - REFERENCE_DENSITY) - traj = create_mdtraj_trajectory_from_simulation_state( simulation_state, - self.data_input_dir / self.name / WATERBOX_N500, + self.data_dir / WATERBOX_N500, ) oxygen_indices = traj.top.select("symbol == O") @@ -262,11 +212,7 @@ def analyze(self) -> WaterRadialDistributionResult: score = (peak_deviation_score + rmse_score) / 2 - # TODO: Remove `range_of_interest=SOLVENT_PEAK_RANGE`? return WaterRadialDistributionResult( - densities=densities, - average_density=average_density, - density_deviation=density_deviation, radii=radii.tolist(), rdf=rdf, mae=mae, @@ -277,34 +223,11 @@ def analyze(self) -> WaterRadialDistributionResult: score=score, ) - @functools.cached_property - def _md_kwargs(self) -> dict[str, Any]: - if self.run_mode == RunMode.DEV: - return SIMULATION_CONFIG_DEV - if self.run_mode == RunMode.FAST: - return SIMULATION_CONFIG_FAST - - return SIMULATION_CONFIG - - @functools.cached_property - def _water_box_n500(self) -> Atoms: - atoms = ase_read(self.data_input_dir / self.name / WATERBOX_N500) - atoms.info["charge"] = DEFAULT_CHARGE - atoms.info["spin"] = DEFAULT_SPIN - return atoms - - @functools.cached_property - def _molecule_indices(self) -> np.ndarray: - molecule_indices = np.load( - self.data_input_dir / self.name / MOLECULE_INDICES_PATH - ) - return molecule_indices - - @functools.cached_property + @property def _reference_data(self): """The experimental reference data for the water RDF benchmark. Contains keys 'r_OO' and 'g_OO', the radii and RDF values. The radii are in Angstrom. """ - return np.load(self.data_input_dir / self.name / REFERENCE_DATA) + return np.load(self.data_dir / REFERENCE_DATA) diff --git a/src/mlipaudit/ui/__init__.py b/src/mlipaudit/ui/__init__.py index b50c93e6..cecad142 100644 --- a/src/mlipaudit/ui/__init__.py +++ b/src/mlipaudit/ui/__init__.py @@ -24,7 +24,9 @@ from mlipaudit.ui.ring_planarity import ring_planarity_page from mlipaudit.ui.sampling import sampling_page from mlipaudit.ui.scaling import scaling_page +from mlipaudit.ui.solvent_density import solvent_density_page from mlipaudit.ui.solvent_radial_distribution import solvent_radial_distribution_page from mlipaudit.ui.stability import stability_page from mlipaudit.ui.tautomers import tautomers_page +from mlipaudit.ui.water_density import water_density_page from mlipaudit.ui.water_radial_distribution import water_radial_distribution_page diff --git a/src/mlipaudit/ui/solvent_density.py b/src/mlipaudit/ui/solvent_density.py new file mode 100644 index 00000000..26d63c56 --- /dev/null +++ b/src/mlipaudit/ui/solvent_density.py @@ -0,0 +1,182 @@ +# 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. + +from typing import Callable, TypeAlias + +import altair as alt +import pandas as pd +import streamlit as st + +from mlipaudit.benchmarks import SolventDensityBenchmark, SolventDensityResult +from mlipaudit.ui.page_wrapper import UIPageWrapper +from mlipaudit.ui.utils import ( + display_failed_models, + display_model_scores, + fetch_selected_models, + filter_failed_results, + get_failed_models, +) + +ModelName: TypeAlias = str +BenchmarkResultForMultipleModels: TypeAlias = dict[ModelName, SolventDensityResult] + + +def _process_data_into_dataframe( + data: BenchmarkResultForMultipleModels, + selected_models: list[str], +) -> pd.DataFrame: + converted_data_scores = [] + for model_name, result in data.items(): + if model_name in selected_models: + model_data_converted = { + "Model name": model_name, + "Score": result.score, + "Average density deviation (g/cm3)": result.avg_density_deviation, + } + for structure_res in result.structures: + if structure_res.failed: + continue + + model_data_converted[ + f"{structure_res.structure_name} density (g/cm3)" + ] = structure_res.average_density + + model_data_converted[ + f"{structure_res.structure_name} density deviation (g/cm3)" + ] = structure_res.density_deviation + converted_data_scores.append(model_data_converted) + df = pd.DataFrame(converted_data_scores) + return df + + +def solvent_density_page( + data_func: Callable[[], BenchmarkResultForMultipleModels], +) -> None: + """Page for the visualization app for the solvent density benchmark. + + Args: + data_func: A data function that delivers the results on request. It does + not take any arguments and returns a dictionary with model names as + keys and the benchmark results objects as values. + """ + st.markdown("# Solvent density") + + st.markdown( + "Here we show the equilibrium density of the solvents CCl4, methanol, and " + "acetonitrile, obtained from NPT simulations. The dashed lines show the " + "reference density of each solvent. A box that expands or collapses will show " + "up as a large deviation from the reference density." + ) + + st.markdown( + "For more information, see the [docs](https://instadeepai.github.io/mlipaudit/" + "benchmarks/molecular_liquids/density.html)." + ) + + if "solvent_density_cached_data" not in st.session_state: + st.session_state.solvent_density_cached_data = data_func() + + data: BenchmarkResultForMultipleModels = ( + st.session_state.solvent_density_cached_data + ) + + if not data: + st.markdown("**No results to display**.") + return + + selected_models = fetch_selected_models(available_models=list(data.keys())) + + if not selected_models: + st.markdown("**No results to display**.") + return + + failed_models = get_failed_models(data) + display_failed_models(failed_models) + data = filter_failed_results(data) + + st.markdown("## Summary statistics") + + df = _process_data_into_dataframe(data, selected_models) + df.sort_values("Score", ascending=False, inplace=True) + display_model_scores(df) + + st.markdown("## Density time series") + + for solvent_index, solvent in enumerate(["CCl4", "methanol", "acetonitrile"]): + plot_data_solvent = [] + reference_density = None + + for model_name, result in data.items(): + if model_name not in selected_models: + continue + if solvent not in result.structure_names: + continue + structure_res = result.structures[solvent_index] + if structure_res.failed or structure_res.densities is None: + continue + + reference_density = structure_res.reference_density + for frame, density in enumerate(structure_res.densities): + plot_data_solvent.append({ + "Frame": frame, + "Density (g/cm3)": density, + "model": str(model_name), + }) + + if not plot_data_solvent: + st.warning(f"No data found for {solvent}") + continue + + st.subheader(f"Density of {solvent}") + + df_plot_solvent = pd.DataFrame(plot_data_solvent) + + chart_solvent = ( + alt.Chart(df_plot_solvent) + .mark_line(strokeWidth=2.0) + .encode( + x=alt.X("Frame:Q", title="Frame"), + y=alt.Y( + "Density (g/cm3):Q", + title="Density (g/cm³)", + scale=alt.Scale(zero=False), + ), + color=alt.Color("model:N", title="Model"), + ) + .properties(width=800, height=400) + ) + + if reference_density is not None: + reference_line = ( + alt.Chart(pd.DataFrame({"y": [reference_density]})) + .mark_rule(color="black", strokeDash=[6, 4], strokeWidth=2) + .encode(y="y:Q") + ) + chart_solvent = chart_solvent + reference_line + + st.altair_chart(chart_solvent, use_container_width=True) + + +class SolventDensityPageWrapper(UIPageWrapper): + """Page wrapper for the solvent density benchmark.""" + + @classmethod + def get_page_func( # noqa: D102 + cls, + ) -> Callable[[Callable[[], BenchmarkResultForMultipleModels]], None]: + return solvent_density_page + + @classmethod + def get_benchmark_class(cls) -> type[SolventDensityBenchmark]: # noqa: D102 + return SolventDensityBenchmark diff --git a/src/mlipaudit/ui/solvent_radial_distribution.py b/src/mlipaudit/ui/solvent_radial_distribution.py index e1bf07bc..4a52ff4d 100644 --- a/src/mlipaudit/ui/solvent_radial_distribution.py +++ b/src/mlipaudit/ui/solvent_radial_distribution.py @@ -57,7 +57,6 @@ def _process_data_into_dataframe( "Model name": model_name, "Score": result.score, "Average peak deviation (Å)": result.avg_peak_deviation, - "Average density deviation (g/cm3)": result.avg_density_deviation, } for structure_res in result.structures: if structure_res.failed: @@ -66,10 +65,6 @@ def _process_data_into_dataframe( model_data_converted[ f"{structure_res.structure_name} peak deviation (Å)" ] = structure_res.peak_deviation - - model_data_converted[ - f"{structure_res.structure_name} density deviation (g/cm3)" - ] = structure_res.density_deviation converted_data_scores.append(model_data_converted) df = pd.DataFrame(converted_data_scores) return df diff --git a/src/mlipaudit/ui/water_density.py b/src/mlipaudit/ui/water_density.py new file mode 100644 index 00000000..2ec91271 --- /dev/null +++ b/src/mlipaudit/ui/water_density.py @@ -0,0 +1,171 @@ +# 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. + +from typing import Callable, TypeAlias + +import altair as alt +import pandas as pd +import streamlit as st + +from mlipaudit.benchmarks import WaterDensityBenchmark, WaterDensityResult +from mlipaudit.ui.page_wrapper import UIPageWrapper +from mlipaudit.ui.utils import ( + display_failed_models, + display_model_scores, + fetch_selected_models, + filter_failed_results, + get_failed_models, +) + +ModelName: TypeAlias = str +BenchmarkResultForMultipleModels: TypeAlias = dict[ModelName, WaterDensityResult] + + +def _process_data_into_dataframe( + data: BenchmarkResultForMultipleModels, + selected_models: list[str], +) -> pd.DataFrame: + converted_data_scores, model_names = [], [] + for model_name, result in data.items(): + if model_name in selected_models: + converted_data_scores.append({ + "Score": result.score, + "Equilibrium density (g/cm3)": result.average_density, + "Density deviation (g/cm3)": result.density_deviation, + "Reference density (g/cm3)": result.reference_density, + }) + model_names.append(model_name) + + df = pd.DataFrame(converted_data_scores, index=model_names) + return df + + +def water_density_page( + data_func: Callable[[], BenchmarkResultForMultipleModels], +) -> None: + """Page for the visualization app for the water density benchmark. + + Args: + data_func: A data function that delivers the results on request. It does + not take any arguments and returns a dictionary with model names as + keys and the benchmark results objects as values. + """ + st.markdown("# Water density") + + st.markdown( + "The equilibrium density of liquid water is a fundamental property that a " + "good MLIP should reproduce. We run an NPT simulation of a box of water " + "molecules and compare the resulting equilibrium density to the experimental " + "reference. A box that expands or collapses will show up as a large deviation " + "from the reference density." + ) + + st.markdown( + "For more information, see the [docs](https://instadeepai.github.io/mlipaudit" + "/benchmarks/molecular_liquids/density.html)." + ) + + if "water_density_cached_data" not in st.session_state: + st.session_state.water_density_cached_data = data_func() + + data: BenchmarkResultForMultipleModels = st.session_state.water_density_cached_data + + if not data: + st.markdown("**No results to display**.") + return + + selected_models = fetch_selected_models(available_models=list(data.keys())) + + if not selected_models: + st.markdown("**No results to display**.") + return + + failed_models = get_failed_models(data) + display_failed_models(failed_models) + data = filter_failed_results(data) + + st.markdown("## Summary statistics") + + df = _process_data_into_dataframe(data, selected_models) + df = df.rename_axis("Model name") + df.sort_values("Score", ascending=False, inplace=True) + display_model_scores(df) + + st.markdown("## Density time series") + st.markdown( + "Here we show the density of the water box over the course of the simulation. " + "The black dashed line shows the experimental reference density. A " + "well-behaved simulation should equilibrate close to this line." + ) + + plot_data = [] + reference_density = None + for model_name, result in data.items(): + if ( + model_name in selected_models + and not result.failed + and result.densities is not None + ): + reference_density = result.reference_density + for frame, density in enumerate(result.densities): + plot_data.append({ + "Frame": frame, + "Density (g/cm3)": density, + "model": str(model_name), + }) + + if not plot_data: + st.markdown("**No density time series to display**.") + return + + df_plot = pd.DataFrame(plot_data) + + chart = ( + alt.Chart(df_plot) + .mark_line(strokeWidth=2.0) + .encode( + x=alt.X("Frame:Q", title="Frame"), + y=alt.Y( + "Density (g/cm3):Q", + title="Density (g/cm³)", + scale=alt.Scale(zero=False), + ), + color=alt.Color("model:N", title="Model"), + ) + .properties(width=800, height=400) + ) + + if reference_density is not None: + reference_line = ( + alt.Chart(pd.DataFrame({"y": [reference_density]})) + .mark_rule(color="black", strokeDash=[6, 4], strokeWidth=2) + .encode(y="y:Q") + ) + chart = chart + reference_line + + st.altair_chart(chart, use_container_width=True) + + +class WaterDensityPageWrapper(UIPageWrapper): + """Page wrapper for the water density benchmark.""" + + @classmethod + def get_page_func( # noqa: D102 + cls, + ) -> Callable[[Callable[[], BenchmarkResultForMultipleModels]], None]: + return water_density_page + + @classmethod + def get_benchmark_class(cls) -> type[WaterDensityBenchmark]: # noqa: D102 + return WaterDensityBenchmark diff --git a/src/mlipaudit/ui/water_radial_distribution.py b/src/mlipaudit/ui/water_radial_distribution.py index 1e2aeef1..32e8dcc1 100644 --- a/src/mlipaudit/ui/water_radial_distribution.py +++ b/src/mlipaudit/ui/water_radial_distribution.py @@ -44,7 +44,6 @@ ] -# TODO: Update this file with density scores @st.cache_resource def _load_tip3p() -> NpzFile: return np.load(WATER_RADIAL_DISTRIBUTION_DATA_DIR / "tip3p_500ps.npz") @@ -73,8 +72,6 @@ def _process_data_into_dataframe( "MAE (Å)": result.mae, "First solvent peak (Å)": result.first_solvent_peak, "Peak deviation (Å)": result.peak_deviation, - "Equilibrium density (g/cm3)": result.average_density, - "Density deviation (g/cm3)": result.density_deviation, } converted_data_scores.append(model_data_converted) model_names.append(model_name) diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index 2b3620d4..825f06f2 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -12,21 +12,140 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Shared configuration and helpers for the molecular liquids benchmarks. + +The water and solvent radial-distribution and density benchmarks all run the same +NPT simulation of a box of molecules. To avoid running these simulations twice, the +RDF and density benchmarks of a given system group share a `reusable_output_id` and +identical `ModelOutput` field signatures (see `benchmarks_cli._transfer_model_output`). +This module centralizes the simulation configuration, input-data loading and density +calculation so that both benchmarks stay in sync. +""" + +import logging +import os +from pathlib import Path +from typing import Any + import numpy as np -from ase import units +from ase import Atoms, units +from ase.io import read as ase_read from mlip.simulation import SimulationState +from mlip.simulation.enums import MDIntegrator + +from mlipaudit.benchmark import DEFAULT_CHARGE, DEFAULT_SPIN, RunModeAsString +from mlipaudit.run_mode import RunMode +from mlipaudit.utils.simulation import run_simulation + +logger = logging.getLogger("mlipaudit") + + +def _as_run_mode(run_mode: RunMode | RunModeAsString) -> RunMode: + """Coerce a run mode that may be given as a string into a `RunMode`. + + Returns: + The run mode as a `RunMode` enum member. + """ + return run_mode if isinstance(run_mode, RunMode) else RunMode(run_mode) + ANGSTROM3_TO_CM3 = 1e-24 +#: Relative-deviation threshold used to score density benchmarks: the equilibrium +#: density is scored via `compute_metric_score(density_deviation / reference_density, +#: DENSITY_RELATIVE_DEVIATION_THRESHOLD, ALPHA)`. A deviation at or below this fraction +#: of the reference density scores 1.0; larger deviations decay exponentially. +#: TODO: Confirm this threshold with the science team (Marco/Silvia) before release. +DENSITY_RELATIVE_DEVIATION_THRESHOLD = 0.05 + +#: Reusable output IDs shared by the RDF and density benchmarks of each system group. +WATER_REUSABLE_OUTPUT_ID = ("water_molecular_liquid",) +SOLVENT_REUSABLE_OUTPUT_ID = ("solvent_molecular_liquid",) + +# -------------------------------------------------------------------------------------- +# Water +# -------------------------------------------------------------------------------------- + +#: Input-data directory name shared by the water RDF and water density benchmarks. +WATER_DATA_NAME = "water_radial_distribution" +WATERBOX_N500 = "water_box_n500_eq.pdb" +WATER_MOLECULE_INDICES_PATH = "water_box_n500_molecule_indices.npy" + +WATER_MOLECULE_WEIGHT = 18.01528 # g/mol +WATER_ATOMS_PER_MOLECULE = 3 +WATER_REFERENCE_DENSITY = 0.997773 # g/cm3 + +WATER_SIMULATION_CONFIG = { + "num_steps": 500_000, + "snapshot_interval": 500, + "num_episodes": 1000, + "temperature_kelvin": 295.15, + "pressure_bar": 1.01325, +} +WATER_SIMULATION_CONFIG_FAST = { + "num_steps": 250_000, + "snapshot_interval": 250, + "num_episodes": 1000, + "temperature_kelvin": 295.15, + "pressure_bar": 1.01325, +} +WATER_SIMULATION_CONFIG_DEV = { + "num_steps": 5, + "snapshot_interval": 1, + "num_episodes": 1, + "temperature_kelvin": 295.15, + "pressure_bar": 1.01325, +} + +# -------------------------------------------------------------------------------------- +# Solvents +# -------------------------------------------------------------------------------------- + +#: Input-data directory name shared by the solvent RDF and solvent density benchmarks. +SOLVENT_DATA_NAME = "solvent_radial_distribution" +NUM_DEV_SYSTEMS = 1 + +#: Per-solvent molecular weight (g/mol) and number of atoms per molecule. The dict +#: order also defines the canonical order of the solvent systems. +SOLVENT_MOLECULE_CONFIG = { + "CCl4": {"molecule_weight": 153.823, "atoms_per_molecule": 5}, + "methanol": {"molecule_weight": 32.042, "atoms_per_molecule": 6}, + "acetonitrile": {"molecule_weight": 41.053, "atoms_per_molecule": 6}, +} +SOLVENT_REFERENCE_DENSITIES = { # g/cm3 + "CCl4": 1.594, + "acetonitrile": 0.786, + "methanol": 0.791, +} + +SOLVENT_SIMULATION_CONFIG = { + "num_steps": 500_000, + "snapshot_interval": 500, + "num_episodes": 1000, + "temperature_kelvin": 293.15, + "pressure_bar": 1.01325, +} +SOLVENT_SIMULATION_CONFIG_FAST = { + "num_steps": 250_000, + "snapshot_interval": 250, + "num_episodes": 1000, + "temperature_kelvin": 293.15, + "pressure_bar": 1.01325, +} +SOLVENT_SIMULATION_CONFIG_DEV = { + "num_steps": 5, + "snapshot_interval": 1, + "num_episodes": 1, + "temperature_kelvin": 293.15, + "pressure_bar": 1.01325, +} + def compute_densities( simulation_state: SimulationState, molecule_weight: float, atoms_per_molecule: int ) -> np.ndarray: """Compute the density (g/cm3) for each frame of the simulation. - Used by `WaterRadialDistributionFunction` and `SolventRadialDistributionFunction`. - TODO: Should be moved into benchmark file when a density benchmark is created. - Args: simulation_state: The final simulation state. molecule_weight: Molecular weight of each solvent molecule. @@ -43,3 +162,175 @@ def compute_densities( densities = density_numerator / density_denominator return densities + + +def average_equilibrated_density(densities: np.ndarray) -> float: + """Average the density over the final four fifths of the frames. + + The first fifth of the trajectory is treated as equilibration and discarded. + + Args: + densities: Per-frame densities (g/cm3). + + Returns: + The average equilibrated density (g/cm3). + """ + n_frames_equilibration = len(densities) // 5 + return float(np.mean(densities[n_frames_equilibration:])) + + +# -------------------------------------------------------------------------------------- +# Water data loading + simulation +# -------------------------------------------------------------------------------------- + + +def get_water_md_kwargs(run_mode: RunMode | RunModeAsString) -> dict[str, Any]: + """Return the water simulation configuration for the given run mode.""" + run_mode = _as_run_mode(run_mode) + if run_mode == RunMode.DEV: + return WATER_SIMULATION_CONFIG_DEV + if run_mode == RunMode.FAST: + return WATER_SIMULATION_CONFIG_FAST + return WATER_SIMULATION_CONFIG + + +def load_water_box(data_dir: str | os.PathLike) -> Atoms: + """Load the water box structure from the given data directory. + + Returns: + The water box as an `ase.Atoms` object with charge and spin set. + """ + atoms = ase_read(Path(data_dir) / WATERBOX_N500) + atoms.info["charge"] = DEFAULT_CHARGE + atoms.info["spin"] = DEFAULT_SPIN + return atoms + + +def load_water_molecule_indices(data_dir: str | os.PathLike) -> np.ndarray: + """Load the per-molecule atom indices used by the NPT barostat. + + Returns: + The per-molecule atom indices. + """ + return np.load(Path(data_dir) / WATER_MOLECULE_INDICES_PATH) + + +def run_water_npt_simulation( + force_field: Any, data_dir: str | os.PathLike, run_mode: RunMode | RunModeAsString +) -> SimulationState | None: + """Run the water box NPT simulation. + + The MD simulation is performed using the JAX MD engine and starts from the + reference structure. The NPT integrator uses Langevin dynamics with a Monte Carlo + barostat. + + Args: + force_field: The force field to run the simulation with. + data_dir: The directory holding the water input data. + run_mode: The run mode controlling the simulation length. + + Returns: + The final simulation state, or None if the simulation failed. + """ + logger.info("Running water box NPT simulation.") + return run_simulation( + atoms=load_water_box(data_dir), + force_field=force_field, + md_integrator=MDIntegrator.NPT_MC_LANGEVIN, + molecule_indices=load_water_molecule_indices(data_dir), + **get_water_md_kwargs(run_mode), + ) + + +# -------------------------------------------------------------------------------------- +# Solvent data loading + simulation +# -------------------------------------------------------------------------------------- + + +def get_solvent_md_kwargs(run_mode: RunMode | RunModeAsString) -> dict[str, Any]: + """Return the solvent simulation configuration for the given run mode.""" + run_mode = _as_run_mode(run_mode) + if run_mode == RunMode.DEV: + return SOLVENT_SIMULATION_CONFIG_DEV + if run_mode == RunMode.FAST: + return SOLVENT_SIMULATION_CONFIG_FAST + return SOLVENT_SIMULATION_CONFIG + + +def get_solvent_system_names(run_mode: RunMode | RunModeAsString) -> list[str]: + """Return the solvent system names to run for the given run mode. + + Returns: + The solvent system names in canonical order. + """ + run_mode = _as_run_mode(run_mode) + system_names = list(SOLVENT_MOLECULE_CONFIG.keys()) + if run_mode == RunMode.STANDARD: + return system_names + # reduced number of cases for DEV and FAST run mode + return system_names[:NUM_DEV_SYSTEMS] + + +def get_solvent_pdb_file_name(system_name: str) -> str: + """Return the PDB file name for a solvent system.""" + return f"{system_name}_eq.pdb" + + +def get_solvent_molecule_indices_file_name(system_name: str) -> str: + """Return the molecule-indices file name for a solvent system.""" + return f"{system_name}_molecule_indices.npy" + + +def load_solvent_system(data_dir: str | os.PathLike, system_name: str) -> Atoms: + """Load a solvent structure from the given data directory. + + Returns: + The solvent structure as an `ase.Atoms` object with charge and spin set. + """ + atoms = ase_read(Path(data_dir) / get_solvent_pdb_file_name(system_name)) + atoms.info["charge"] = DEFAULT_CHARGE + atoms.info["spin"] = DEFAULT_SPIN + return atoms + + +def load_solvent_molecule_indices( + data_dir: str | os.PathLike, system_name: str +) -> np.ndarray: + """Load the per-molecule atom indices for a solvent system. + + Returns: + The per-molecule atom indices for the given solvent system. + """ + return np.load(Path(data_dir) / get_solvent_molecule_indices_file_name(system_name)) + + +def run_solvent_npt_simulations( + force_field: Any, data_dir: str | os.PathLike, run_mode: RunMode | RunModeAsString +) -> tuple[list[str], list[SimulationState | None]]: + """Run one NPT simulation per solvent system. + + Args: + force_field: The force field to run the simulations with. + data_dir: The directory holding the solvent input data. + run_mode: The run mode controlling the simulation length and system count. + + Returns: + A tuple of (system names, simulation states) in matching order. A simulation + state is None if the corresponding simulation failed. + """ + system_names = get_solvent_system_names(run_mode) + md_kwargs = get_solvent_md_kwargs(run_mode) + + simulation_states: list[SimulationState | None] = [] + for system_name in system_names: + logger.info("Running NPT simulation for %s.", system_name) + simulation_state = run_simulation( + atoms=load_solvent_system(data_dir, system_name), + force_field=force_field, + md_integrator=MDIntegrator.NPT_MC_LANGEVIN, + molecule_indices=load_solvent_molecule_indices(data_dir, system_name), + **md_kwargs, + ) + simulation_states.append(simulation_state) + + return system_names, simulation_states diff --git a/tests/solvent_density/test_solvent_density.py b/tests/solvent_density/test_solvent_density.py new file mode 100644 index 00000000..87e95f9d --- /dev/null +++ b/tests/solvent_density/test_solvent_density.py @@ -0,0 +1,136 @@ +# 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. + +import re +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pytest +from ase.io import read as ase_read +from mlip.simulation import SimulationState + +from mlipaudit.benchmarks import ( + SolventDensityBenchmark, + SolventDensityModelOutput, + SolventDensityResult, + SolventRadialDistributionBenchmark, + SolventRadialDistributionModelOutput, +) +from mlipaudit.benchmarks_cli import _transfer_model_output +from mlipaudit.run_mode import RunMode + +INPUT_DATA_DIR = Path(__file__).parent.parent / "data" +# The solvent density benchmark reuses the solvent RDF benchmark's input data. +SOLVENT_DATA_DIR = INPUT_DATA_DIR / "solvent_radial_distribution" + + +@pytest.fixture +def solvent_density_benchmark( + request, + mocked_benchmark_init, # Use the generic init mock + mock_force_field, # Use the generic force field mock +) -> SolventDensityBenchmark: + """Assembles a fully configured and isolated SolventDensityBenchmark instance. + + Returns: + An initialized SolventDensityBenchmark instance. + """ + is_fast_run = getattr(request, "param", False) + run_mode = RunMode.DEV if is_fast_run else RunMode.STANDARD + + return SolventDensityBenchmark( + force_field=mock_force_field, + data_input_dir=INPUT_DATA_DIR, + run_mode=run_mode, + ) + + +def _mock_ccl4_simulation_state() -> SimulationState: + """Build a stationary CCl4 trajectory from the equilibrated structure.""" + atoms = ase_read(SOLVENT_DATA_DIR / "CCl4_eq.pdb") + num_frames = 2 + positions = np.tile(np.array(atoms.positions), reps=(num_frames, 1, 1)) + cells = np.tile(np.array(atoms.get_cell()), reps=(num_frames, 1, 1)) + return SimulationState( + positions=positions, temperature=np.ones(num_frames), cell=cells + ) + + +def test_data_name_points_to_shared_data(solvent_density_benchmark): + """The density benchmark should read from the RDF benchmark's data directory.""" + assert solvent_density_benchmark.data_name == "solvent_radial_distribution" + assert solvent_density_benchmark.data_dir == SOLVENT_DATA_DIR + + +@pytest.mark.parametrize("solvent_density_benchmark", [True], indirect=True) +def test_full_run_with_mocked_engine( + solvent_density_benchmark, mock_jaxmd_simulation_engine +): + """Integration test testing a full run of the benchmark.""" + benchmark = solvent_density_benchmark + mock_engine = mock_jaxmd_simulation_engine() + with patch( + "mlipaudit.utils.simulation.JaxMDSimulationEngine", + return_value=mock_engine, + ) as mock_engine_class: + benchmark.run_model() + + assert mock_engine_class.call_count == 1 + assert isinstance(benchmark.model_output, SolventDensityModelOutput) + + benchmark.model_output = SolventDensityModelOutput( + structure_names=["CCl4"], + simulation_states=[_mock_ccl4_simulation_state()], + ) + + result = benchmark.analyze() + assert type(result) is SolventDensityResult + + assert len(result.structures) == 1 + assert result.structures[0].structure_name == "CCl4" + + # Target density = 1.594, initial density = 1.368 + assert 1.3 < result.structures[0].average_density < 1.6 + assert result.structures[0].density_deviation < 0.3 + assert 0.0 <= result.score <= 1.0 + + +def test_reuses_solvent_rdf_simulation_output(): + """The density benchmark must be able to reuse the RDF benchmark's model output. + + This mirrors the CLI reuse path in `benchmarks_cli.run_benchmarks`, where a + cached `ModelOutput` is transferred between benchmarks that share a + `reusable_output_id`. + """ + assert ( + SolventDensityBenchmark.reusable_output_id + == SolventRadialDistributionBenchmark.reusable_output_id + ) + + rdf_output = SolventRadialDistributionModelOutput( + structure_names=["CCl4"], + simulation_states=[_mock_ccl4_simulation_state()], + ) + transferred = _transfer_model_output(rdf_output, SolventDensityModelOutput) + + assert isinstance(transferred, SolventDensityModelOutput) + assert transferred.structure_names == ["CCl4"] + + +def test_analyze_raises_error_if_run_first(solvent_density_benchmark): + """Verifies the RuntimeError is raised when analyze is called first.""" + expected_message = "Must call run_model() first." + with pytest.raises(RuntimeError, match=re.escape(expected_message)): + solvent_density_benchmark.analyze() diff --git a/tests/solvent_radial_distribution/test_solvent_radial_distribution.py b/tests/solvent_radial_distribution/test_solvent_radial_distribution.py index e96970a6..3bd2054d 100644 --- a/tests/solvent_radial_distribution/test_solvent_radial_distribution.py +++ b/tests/solvent_radial_distribution/test_solvent_radial_distribution.py @@ -95,10 +95,6 @@ def test_full_run_with_mocked_engine( assert 5.4 < result.structures[0].first_solvent_peak < 6.4 assert result.structures[0].peak_deviation < 0.5 - # Target density = 1.594, initial density = 1.368 - assert 1.3 < result.structures[0].average_density < 1.6 - assert result.structures[0].density_deviation < 0.3 - def test_analyze_raises_error_if_run_first(solvent_radial_distribution_benchmark): """Verifies the RuntimeError using the new fixture.""" diff --git a/tests/test_ui_pages.py b/tests/test_ui_pages.py index 6570927b..64ac35fc 100644 --- a/tests/test_ui_pages.py +++ b/tests/test_ui_pages.py @@ -29,11 +29,16 @@ RingPlanarityBenchmark, SamplingBenchmark, ScalingBenchmark, + SolventDensityBenchmark, + SolventDensityResult, + SolventDensityStructureResult, SolventRadialDistributionBenchmark, SolventRadialDistributionResult, SolventRadialDistributionStructureResult, StabilityBenchmark, TautomersBenchmark, + WaterDensityBenchmark, + WaterDensityResult, WaterRadialDistributionBenchmark, ) from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( @@ -68,9 +73,11 @@ ring_planarity_page, sampling_page, scaling_page, + solvent_density_page, solvent_radial_distribution_page, stability_page, tautomers_page, + water_density_page, water_radial_distribution_page, ) @@ -102,6 +109,7 @@ def _add_failed_molecule( elif benchmark_class in [ FoldingStabilityBenchmark, SamplingBenchmark, + SolventDensityBenchmark, SolventRadialDistributionBenchmark, StabilityBenchmark, ScalingBenchmark, @@ -214,6 +222,19 @@ def _add_failed_model(benchmark_class, model_results) -> dict[str, BenchmarkResu ) elif benchmark_class is WaterRadialDistributionBenchmark: model_results["model_3"] = WaterRadialDistributionResult(failed=True, score=0.0) + elif benchmark_class is SolventDensityBenchmark: + model_results["model_3"] = SolventDensityResult( + structure_names=["failed_mol"], + structures=[ + SolventDensityStructureResult( + structure_name="failed_mol", failed=True, score=0.0 + ) + ], + failed=True, + score=0.0, + ) + elif benchmark_class is WaterDensityBenchmark: + model_results["model_3"] = WaterDensityResult(failed=True, score=0.0) return model_results @@ -309,6 +330,7 @@ def data_func() -> BenchmarkResultForMultipleModels: # Manually add the score for the test if benchmark_class not in [ ScalingBenchmark, + SolventDensityBenchmark, SolventRadialDistributionBenchmark, ]: kwargs_for_result["score"] = 0.3 @@ -369,9 +391,11 @@ def _app_script(page_func, data_func, scores, is_public): (NoncovalentInteractionsBenchmark, noncovalent_interactions_page), (ReferenceGeometryStabilityBenchmark, reference_geometry_stability_page), (SolventRadialDistributionBenchmark, solvent_radial_distribution_page), + (SolventDensityBenchmark, solvent_density_page), (StabilityBenchmark, stability_page), (TautomersBenchmark, tautomers_page), (WaterRadialDistributionBenchmark, water_radial_distribution_page), + (WaterDensityBenchmark, water_density_page), (ScalingBenchmark, scaling_page), (SamplingBenchmark, sampling_page), (DihedralScanBenchmark, dihedral_scan_page), diff --git a/tests/water_density/test_water_density.py b/tests/water_density/test_water_density.py new file mode 100644 index 00000000..13797b6c --- /dev/null +++ b/tests/water_density/test_water_density.py @@ -0,0 +1,133 @@ +# 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. + +import re +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pytest +from mlip.simulation import SimulationState + +from mlipaudit.benchmarks import ( + WaterDensityBenchmark, + WaterDensityModelOutput, + WaterDensityResult, + WaterRadialDistributionBenchmark, + WaterRadialDistributionModelOutput, +) +from mlipaudit.benchmarks_cli import _transfer_model_output +from mlipaudit.run_mode import RunMode + +INPUT_DATA_DIR = Path(__file__).parent.parent / "data" +# The water density benchmark reuses the water RDF benchmark's input data. +WATER_DATA_DIR = INPUT_DATA_DIR / "water_radial_distribution" + + +@pytest.fixture +def water_density_benchmark( + request, + mocked_benchmark_init, # Use the generic init mock + mock_force_field, # Use the generic force field mock +) -> WaterDensityBenchmark: + """Assembles a fully configured and isolated WaterDensityBenchmark instance. + + Returns: + An initialized WaterDensityBenchmark instance. + """ + is_fast_run = getattr(request, "param", False) + run_mode = RunMode.DEV if is_fast_run else RunMode.STANDARD + + return WaterDensityBenchmark( + force_field=mock_force_field, + data_input_dir=INPUT_DATA_DIR, + run_mode=run_mode, + ) + + +def _mock_simulation_state() -> SimulationState: + """Build a stationary water box trajectory with a fixed 24.772 Å box.""" + num_frames = 2 + positions = np.tile( + np.load(WATER_DATA_DIR / "positions.npy"), + reps=(num_frames, 1, 1), + ) + cells = np.tile(24.772 * np.eye(3), reps=(num_frames, 1, 1)) + return SimulationState( + positions=positions, temperature=np.ones(num_frames), cell=cells + ) + + +def test_data_name_points_to_shared_data(water_density_benchmark): + """The density benchmark should read from the RDF benchmark's data directory.""" + assert water_density_benchmark.data_name == "water_radial_distribution" + assert water_density_benchmark.data_dir == WATER_DATA_DIR + + +@pytest.mark.parametrize("water_density_benchmark", [True, False], indirect=True) +def test_full_run_with_mocked_engine( + water_density_benchmark, mock_jaxmd_simulation_engine +): + """Integration test testing a full run of the benchmark.""" + benchmark = water_density_benchmark + mock_engine = mock_jaxmd_simulation_engine() + with patch( + "mlipaudit.utils.simulation.JaxMDSimulationEngine", + return_value=mock_engine, + ) as mock_engine_class: + benchmark.run_model() + + assert mock_engine_class.call_count == 1 + assert isinstance(benchmark.model_output, WaterDensityModelOutput) + + benchmark.model_output = WaterDensityModelOutput( + simulation_state=_mock_simulation_state() + ) + + result = benchmark.analyze() + assert type(result) is WaterDensityResult + + # Target density = 0.997773, initial density = 0.9859266 + assert 0.9 < result.average_density < 1.1 + assert result.density_deviation < 0.1 + assert 0.0 <= result.score <= 1.0 + + +def test_reuses_water_rdf_simulation_output(): + """The density benchmark must be able to reuse the RDF benchmark's model output. + + This mirrors the CLI reuse path in `benchmarks_cli.run_benchmarks`, where a + cached `ModelOutput` is transferred between benchmarks that share a + `reusable_output_id`. + """ + assert ( + WaterDensityBenchmark.reusable_output_id + == WaterRadialDistributionBenchmark.reusable_output_id + ) + + rdf_output = WaterRadialDistributionModelOutput( + simulation_state=_mock_simulation_state(), failed=False + ) + transferred = _transfer_model_output(rdf_output, WaterDensityModelOutput) + + assert isinstance(transferred, WaterDensityModelOutput) + assert transferred.simulation_state is rdf_output.simulation_state + assert transferred.failed is False + + +def test_analyze_raises_error_if_run_first(water_density_benchmark): + """Verifies the RuntimeError is raised when analyze is called first.""" + expected_message = "Must call run_model() first." + with pytest.raises(RuntimeError, match=re.escape(expected_message)): + water_density_benchmark.analyze() diff --git a/tests/water_radial_distribution/test_water_radial_distribution.py b/tests/water_radial_distribution/test_water_radial_distribution.py index 37eda854..f3a9e473 100644 --- a/tests/water_radial_distribution/test_water_radial_distribution.py +++ b/tests/water_radial_distribution/test_water_radial_distribution.py @@ -97,10 +97,6 @@ def test_full_run_with_mocked_engine( max_radii = np.array(result.radii)[np.argmax(np.array(result.rdf))] assert 2.5 < max_radii < 3.0 - # Target density = 0.997773, initial density = 0.9859266 - assert 0.9 < result.average_density < 1.1 - assert result.density_deviation < 0.1 - def test_analyze_raises_error_if_run_first(water_radial_distribution_benchmark): """Verifies the RuntimeError using the new fixture.""" From 511261bdebd8d122a1f10c3581ebacb6843f8965 Mon Sep 17 00:00:00 2001 From: lwalew Date: Fri, 3 Jul 2026 16:19:08 +0200 Subject: [PATCH 02/12] refactor: type Benchmark.run_mode as RunMode, drop _as_run_mode helper `Benchmark.__init__` already coerces `run_mode` to a `RunMode`, so annotate the stored attribute as `RunMode` and let the molecular-liquids helpers take a plain `RunMode`. Removes the unnecessary `_as_run_mode` normalization. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mlipaudit/benchmark.py | 6 +++--- src/mlipaudit/utils/molecular_liquids.py | 25 ++++++------------------ 2 files changed, 9 insertions(+), 22 deletions(-) diff --git a/src/mlipaudit/benchmark.py b/src/mlipaudit/benchmark.py index c94a355b..0afa836a 100644 --- a/src/mlipaudit/benchmark.py +++ b/src/mlipaudit/benchmark.py @@ -140,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 diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index 825f06f2..658bc389 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -33,22 +33,12 @@ from mlip.simulation import SimulationState from mlip.simulation.enums import MDIntegrator -from mlipaudit.benchmark import DEFAULT_CHARGE, DEFAULT_SPIN, RunModeAsString +from mlipaudit.benchmark import DEFAULT_CHARGE, DEFAULT_SPIN from mlipaudit.run_mode import RunMode from mlipaudit.utils.simulation import run_simulation logger = logging.getLogger("mlipaudit") - -def _as_run_mode(run_mode: RunMode | RunModeAsString) -> RunMode: - """Coerce a run mode that may be given as a string into a `RunMode`. - - Returns: - The run mode as a `RunMode` enum member. - """ - return run_mode if isinstance(run_mode, RunMode) else RunMode(run_mode) - - ANGSTROM3_TO_CM3 = 1e-24 #: Relative-deviation threshold used to score density benchmarks: the equilibrium @@ -184,9 +174,8 @@ def average_equilibrated_density(densities: np.ndarray) -> float: # -------------------------------------------------------------------------------------- -def get_water_md_kwargs(run_mode: RunMode | RunModeAsString) -> dict[str, Any]: +def get_water_md_kwargs(run_mode: RunMode) -> dict[str, Any]: """Return the water simulation configuration for the given run mode.""" - run_mode = _as_run_mode(run_mode) if run_mode == RunMode.DEV: return WATER_SIMULATION_CONFIG_DEV if run_mode == RunMode.FAST: @@ -216,7 +205,7 @@ def load_water_molecule_indices(data_dir: str | os.PathLike) -> np.ndarray: def run_water_npt_simulation( - force_field: Any, data_dir: str | os.PathLike, run_mode: RunMode | RunModeAsString + force_field: Any, data_dir: str | os.PathLike, run_mode: RunMode ) -> SimulationState | None: """Run the water box NPT simulation. @@ -247,9 +236,8 @@ def run_water_npt_simulation( # -------------------------------------------------------------------------------------- -def get_solvent_md_kwargs(run_mode: RunMode | RunModeAsString) -> dict[str, Any]: +def get_solvent_md_kwargs(run_mode: RunMode) -> dict[str, Any]: """Return the solvent simulation configuration for the given run mode.""" - run_mode = _as_run_mode(run_mode) if run_mode == RunMode.DEV: return SOLVENT_SIMULATION_CONFIG_DEV if run_mode == RunMode.FAST: @@ -257,13 +245,12 @@ def get_solvent_md_kwargs(run_mode: RunMode | RunModeAsString) -> dict[str, Any] return SOLVENT_SIMULATION_CONFIG -def get_solvent_system_names(run_mode: RunMode | RunModeAsString) -> list[str]: +def get_solvent_system_names(run_mode: RunMode) -> list[str]: """Return the solvent system names to run for the given run mode. Returns: The solvent system names in canonical order. """ - run_mode = _as_run_mode(run_mode) system_names = list(SOLVENT_MOLECULE_CONFIG.keys()) if run_mode == RunMode.STANDARD: return system_names @@ -305,7 +292,7 @@ def load_solvent_molecule_indices( def run_solvent_npt_simulations( - force_field: Any, data_dir: str | os.PathLike, run_mode: RunMode | RunModeAsString + force_field: Any, data_dir: str | os.PathLike, run_mode: RunMode ) -> tuple[list[str], list[SimulationState | None]]: """Run one NPT simulation per solvent system. From da3481d7755641d81045a241c46a782de5e7683a Mon Sep 17 00:00:00 2001 From: lwalew Date: Fri, 3 Jul 2026 16:25:11 +0200 Subject: [PATCH 03/12] fix: use Leon's density scoring params (2% threshold, alpha=0.1) Match the density score to Leon's NPT density analysis: relative-deviation threshold of 0.02 and a gentle decay (DENSITY_SCORE_ALPHA=0.1), rather than the placeholder 0.05 with the steep global scoring.ALPHA=3.0. With these values the density scores reproduce the expected ranking (v1 models best on water, v2 models poor). The ideal per-solvent threshold should ultimately be based on the isothermal compressibility; still pending final sign-off. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../benchmarks/molecular_liquids/density.rst | 8 +++++--- .../solvent_density/solvent_density.py | 5 +++-- .../benchmarks/water_density/water_density.py | 5 +++-- src/mlipaudit/utils/molecular_liquids.py | 17 +++++++++++------ 4 files changed, 22 insertions(+), 13 deletions(-) diff --git a/docs/source/benchmarks/molecular_liquids/density.rst b/docs/source/benchmarks/molecular_liquids/density.rst index 836138da..0ebd921f 100644 --- a/docs/source/benchmarks/molecular_liquids/density.rst +++ b/docs/source/benchmarks/molecular_liquids/density.rst @@ -52,6 +52,8 @@ 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 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. +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. diff --git a/src/mlipaudit/benchmarks/solvent_density/solvent_density.py b/src/mlipaudit/benchmarks/solvent_density/solvent_density.py index 16ae23fa..bc88e6b5 100644 --- a/src/mlipaudit/benchmarks/solvent_density/solvent_density.py +++ b/src/mlipaudit/benchmarks/solvent_density/solvent_density.py @@ -23,9 +23,10 @@ BenchmarkResult, ModelOutput, ) -from mlipaudit.scoring import ALPHA, compute_metric_score +from mlipaudit.scoring import compute_metric_score from mlipaudit.utils.molecular_liquids import ( DENSITY_RELATIVE_DEVIATION_THRESHOLD, + DENSITY_SCORE_ALPHA, SOLVENT_DATA_NAME, SOLVENT_MOLECULE_CONFIG, SOLVENT_REFERENCE_DENSITIES, @@ -184,7 +185,7 @@ def analyze(self) -> SolventDensityResult: score = compute_metric_score( np.array([relative_deviation]), DENSITY_RELATIVE_DEVIATION_THRESHOLD, - ALPHA, + DENSITY_SCORE_ALPHA, ).item() structure_results.append( diff --git a/src/mlipaudit/benchmarks/water_density/water_density.py b/src/mlipaudit/benchmarks/water_density/water_density.py index cff2e988..9d09a16e 100644 --- a/src/mlipaudit/benchmarks/water_density/water_density.py +++ b/src/mlipaudit/benchmarks/water_density/water_density.py @@ -22,9 +22,10 @@ BenchmarkResult, ModelOutput, ) -from mlipaudit.scoring import ALPHA, compute_metric_score +from mlipaudit.scoring import compute_metric_score from mlipaudit.utils.molecular_liquids import ( DENSITY_RELATIVE_DEVIATION_THRESHOLD, + DENSITY_SCORE_ALPHA, WATER_ATOMS_PER_MOLECULE, WATER_DATA_NAME, WATER_MOLECULE_WEIGHT, @@ -146,7 +147,7 @@ def analyze(self) -> WaterDensityResult: score = compute_metric_score( np.array([relative_deviation]), DENSITY_RELATIVE_DEVIATION_THRESHOLD, - ALPHA, + DENSITY_SCORE_ALPHA, ).item() return WaterDensityResult( diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index 658bc389..f2392895 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -41,12 +41,17 @@ ANGSTROM3_TO_CM3 = 1e-24 -#: Relative-deviation threshold used to score density benchmarks: the equilibrium -#: density is scored via `compute_metric_score(density_deviation / reference_density, -#: DENSITY_RELATIVE_DEVIATION_THRESHOLD, ALPHA)`. A deviation at or below this fraction -#: of the reference density scores 1.0; larger deviations decay exponentially. -#: TODO: Confirm this threshold with the science team (Marco/Silvia) before release. -DENSITY_RELATIVE_DEVIATION_THRESHOLD = 0.05 +#: Scoring parameters for the density benchmarks: the equilibrium density is scored +#: via `compute_metric_score(density_deviation / reference_density, +#: DENSITY_RELATIVE_DEVIATION_THRESHOLD, DENSITY_SCORE_ALPHA)`. A relative deviation at +#: or below the threshold scores 1.0; larger deviations decay exponentially at a rate +#: set by alpha. These follow Leon's NPT density analysis (2% target, gentle decay); +#: the ideal threshold should ultimately be based on the isothermal compressibility of +#: each solvent. Note the density benchmarks use their own alpha rather than the global +#: `scoring.ALPHA` (which is much steeper). +#: TODO: Confirm these with the science team (Leon/Marco/Silvia) before release. +DENSITY_RELATIVE_DEVIATION_THRESHOLD = 0.02 +DENSITY_SCORE_ALPHA = 0.1 #: Reusable output IDs shared by the RDF and density benchmarks of each system group. WATER_REUSABLE_OUTPUT_ID = ("water_molecular_liquid",) From ce664a6b73bc80ad9f768321797dc567deffce07 Mon Sep 17 00:00:00 2001 From: lwalew Date: Fri, 3 Jul 2026 16:26:19 +0200 Subject: [PATCH 04/12] chore: drop verbose comment on density scoring constants Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mlipaudit/utils/molecular_liquids.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index f2392895..0474ae6a 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -41,15 +41,6 @@ ANGSTROM3_TO_CM3 = 1e-24 -#: Scoring parameters for the density benchmarks: the equilibrium density is scored -#: via `compute_metric_score(density_deviation / reference_density, -#: DENSITY_RELATIVE_DEVIATION_THRESHOLD, DENSITY_SCORE_ALPHA)`. A relative deviation at -#: or below the threshold scores 1.0; larger deviations decay exponentially at a rate -#: set by alpha. These follow Leon's NPT density analysis (2% target, gentle decay); -#: the ideal threshold should ultimately be based on the isothermal compressibility of -#: each solvent. Note the density benchmarks use their own alpha rather than the global -#: `scoring.ALPHA` (which is much steeper). -#: TODO: Confirm these with the science team (Leon/Marco/Silvia) before release. DENSITY_RELATIVE_DEVIATION_THRESHOLD = 0.02 DENSITY_SCORE_ALPHA = 0.1 From d374e801723c433499ae32793c4bb6280cf49936 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 15:07:36 +0200 Subject: [PATCH 05/12] refactor: look up solvent structures by name in the UI, drop order comment Address review feedback that solvent dict order isn't meaningful. The solvent UI pages no longer hardcode ["CCl4", "methanol", "acetonitrile"] or index `structures` positionally; they derive the system list from the results and look structures up by name (new `ordered_structure_names` helper). Also drop the misleading "canonical order" comment/docstring, since scoring is entirely by-name and nothing depends on the config dict ordering. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mlipaudit/ui/solvent_density.py | 23 +++++--- .../ui/solvent_radial_distribution.py | 57 ++++++++++--------- src/mlipaudit/ui/utils.py | 25 ++++++++ src/mlipaudit/utils/molecular_liquids.py | 19 ++----- 4 files changed, 73 insertions(+), 51 deletions(-) diff --git a/src/mlipaudit/ui/solvent_density.py b/src/mlipaudit/ui/solvent_density.py index 26d63c56..32978337 100644 --- a/src/mlipaudit/ui/solvent_density.py +++ b/src/mlipaudit/ui/solvent_density.py @@ -26,6 +26,7 @@ fetch_selected_models, filter_failed_results, get_failed_models, + ordered_structure_names, ) ModelName: TypeAlias = str @@ -73,10 +74,10 @@ def solvent_density_page( st.markdown("# Solvent density") st.markdown( - "Here we show the equilibrium density of the solvents CCl4, methanol, and " - "acetonitrile, obtained from NPT simulations. The dashed lines show the " - "reference density of each solvent. A box that expands or collapses will show " - "up as a large deviation from the reference density." + "Here we show the equilibrium density of each molecular solvent, obtained from " + "NPT simulations. The dashed lines show the reference density of each solvent. " + "A box that expands or collapses will show up as a large deviation from the " + "reference density." ) st.markdown( @@ -113,17 +114,21 @@ def solvent_density_page( st.markdown("## Density time series") - for solvent_index, solvent in enumerate(["CCl4", "methanol", "acetonitrile"]): + for solvent in ordered_structure_names(data, selected_models): plot_data_solvent = [] reference_density = None for model_name, result in data.items(): if model_name not in selected_models: continue - if solvent not in result.structure_names: - continue - structure_res = result.structures[solvent_index] - if structure_res.failed or structure_res.densities is None: + structure_res = {s.structure_name: s for s in result.structures}.get( + solvent + ) + if ( + structure_res is None + or structure_res.failed + or structure_res.densities is None + ): continue reference_density = structure_res.reference_density diff --git a/src/mlipaudit/ui/solvent_radial_distribution.py b/src/mlipaudit/ui/solvent_radial_distribution.py index 4a52ff4d..a688dcbf 100644 --- a/src/mlipaudit/ui/solvent_radial_distribution.py +++ b/src/mlipaudit/ui/solvent_radial_distribution.py @@ -33,6 +33,7 @@ fetch_selected_models, filter_failed_results, get_failed_models, + ordered_structure_names, ) APP_DATA_DIR = Path(__file__).parent.parent / "app_data" @@ -94,9 +95,9 @@ def solvent_radial_distribution_page( st.markdown("# Solvent Radial distribution function") st.markdown( - "Here we show the radial distribution function of the solvents CCl4, " - "methanol, and acetonitrile. The vertical lines show the reference " - "maximum of the radial distribution function for each solvent." + "Here we show the radial distribution function of each solvent. The vertical " + "lines show the reference maximum of the radial distribution function for each " + "solvent." ) st.markdown( @@ -138,25 +139,28 @@ def solvent_radial_distribution_page( st.markdown("## Radial distribution functions") - for solvent_index, solvent in enumerate(["CCl4", "methanol", "acetonitrile"]): + for solvent in ordered_structure_names(data, selected_models): rdf_data_solvent = {} for model_name, result in data.items(): - if ( - model_name in selected_models - and solvent in result.structure_names - and not result.structures[solvent_index].failed - ): - rdf_data_solvent[model_name] = { - "r": np.array(result.structures[solvent_index].radii), - "rdf": np.array(result.structures[solvent_index].rdf), - } + if model_name not in selected_models: + continue + structure_res = {s.structure_name: s for s in result.structures}.get( + solvent + ) + if structure_res is None or structure_res.failed: + continue + rdf_data_solvent[model_name] = { + "r": np.array(structure_res.radii), + "rdf": np.array(structure_res.rdf), + } if len(rdf_data_solvent) > 0: - st.subheader( - f"Radial distribution function of {solvent} " - f"({solvent_maxima[solvent]['type']})" - ) + maximum = solvent_maxima.get(solvent) + title = f"Radial distribution function of {solvent}" + if maximum is not None: + title += f" ({maximum['type']})" + st.subheader(title) # Convert to long format for Altair plotting plot_data_solvent = [] @@ -186,17 +190,16 @@ def solvent_radial_distribution_page( .properties(width=800, height=400) ) - # Add vertical line at experimental maximum - vline = ( - alt.Chart(pd.DataFrame({"x": [solvent_maxima[solvent]["distance"]]})) - .mark_rule(color="black", strokeWidth=2) - .encode(x="x:Q") - ) - - # Combine the line chart with the vertical line - combined_chart = chart_solvent + vline + # Add vertical line at experimental maximum, if available + if maximum is not None: + vline = ( + alt.Chart(pd.DataFrame({"x": [maximum["distance"]]})) + .mark_rule(color="black", strokeWidth=2) + .encode(x="x:Q") + ) + chart_solvent = chart_solvent + vline - st.altair_chart(combined_chart, use_container_width=True) + st.altair_chart(chart_solvent, use_container_width=True) else: st.warning(f"No data found for {solvent}") diff --git a/src/mlipaudit/ui/utils.py b/src/mlipaudit/ui/utils.py index 55cfdd50..3c9ede38 100644 --- a/src/mlipaudit/ui/utils.py +++ b/src/mlipaudit/ui/utils.py @@ -326,6 +326,31 @@ def display_failed_models(model_names: list[str]) -> None: st.markdown("Models that failed to run: \n" + markdown_list) +def ordered_structure_names( + data: BenchmarkResultForMultipleModels, selected_models: list[str] +) -> list[str]: + """Collect the de-duplicated structure names across the selected models' results. + + Systems are identified by name rather than by position, so the UI does not depend + on the order in which they were simulated. Names are returned in first-seen order. + + Args: + data: The dictionary of results for a given benchmark. + selected_models: The models to consider. + + Returns: + The structure names in first-seen order. + """ + names: list[str] = [] + for model_name, result in data.items(): + if model_name not in selected_models: + continue + for name in getattr(result, "structure_names", []): + if name not in names: + names.append(name) + return names + + def filter_failed_results( data: BenchmarkResultForMultipleModels, ) -> BenchmarkResultForMultipleModels: diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index 0474ae6a..d19843d1 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -12,16 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared configuration and helpers for the molecular liquids benchmarks. - -The water and solvent radial-distribution and density benchmarks all run the same -NPT simulation of a box of molecules. To avoid running these simulations twice, the -RDF and density benchmarks of a given system group share a `reusable_output_id` and -identical `ModelOutput` field signatures (see `benchmarks_cli._transfer_model_output`). -This module centralizes the simulation configuration, input-data loading and density -calculation so that both benchmarks stay in sync. -""" - import logging import os from pathlib import Path @@ -91,8 +81,7 @@ SOLVENT_DATA_NAME = "solvent_radial_distribution" NUM_DEV_SYSTEMS = 1 -#: Per-solvent molecular weight (g/mol) and number of atoms per molecule. The dict -#: order also defines the canonical order of the solvent systems. +#: Per-solvent molecular weight (g/mol) and number of atoms per molecule. SOLVENT_MOLECULE_CONFIG = { "CCl4": {"molecule_weight": 153.823, "atoms_per_molecule": 5}, "methanol": {"molecule_weight": 32.042, "atoms_per_molecule": 6}, @@ -179,7 +168,7 @@ def get_water_md_kwargs(run_mode: RunMode) -> dict[str, Any]: return WATER_SIMULATION_CONFIG -def load_water_box(data_dir: str | os.PathLike) -> Atoms: +def load_water_system(data_dir: str | os.PathLike) -> Atoms: """Load the water box structure from the given data directory. Returns: @@ -219,7 +208,7 @@ def run_water_npt_simulation( """ logger.info("Running water box NPT simulation.") return run_simulation( - atoms=load_water_box(data_dir), + atoms=load_water_system(data_dir), force_field=force_field, md_integrator=MDIntegrator.NPT_MC_LANGEVIN, molecule_indices=load_water_molecule_indices(data_dir), @@ -245,7 +234,7 @@ def get_solvent_system_names(run_mode: RunMode) -> list[str]: """Return the solvent system names to run for the given run mode. Returns: - The solvent system names in canonical order. + The solvent system names. """ system_names = list(SOLVENT_MOLECULE_CONFIG.keys()) if run_mode == RunMode.STANDARD: From 49cee31832a4d5e898a58795f071853ae71b58b0 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 15:28:04 +0200 Subject: [PATCH 06/12] feat: address comments Docs and make some functions private --- .../benchmarks/molecular_liquids/density.rst | 2 +- src/mlipaudit/utils/molecular_liquids.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/source/benchmarks/molecular_liquids/density.rst b/docs/source/benchmarks/molecular_liquids/density.rst index 0ebd921f..9434ebc5 100644 --- a/docs/source/benchmarks/molecular_liquids/density.rst +++ b/docs/source/benchmarks/molecular_liquids/density.rst @@ -20,7 +20,7 @@ The benchmark runs the same **MD** simulation as the :ref:`radial_distribution` **NPT** simulation using the **MLIP** model for **500,000 steps**, leveraging the `jax-md `_ engine from the `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 RDF and density +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. diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index d19843d1..e3cecff6 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -159,7 +159,7 @@ def average_equilibrated_density(densities: np.ndarray) -> float: # -------------------------------------------------------------------------------------- -def get_water_md_kwargs(run_mode: RunMode) -> dict[str, Any]: +def _get_water_md_kwargs(run_mode: RunMode) -> dict[str, Any]: """Return the water simulation configuration for the given run mode.""" if run_mode == RunMode.DEV: return WATER_SIMULATION_CONFIG_DEV @@ -180,7 +180,7 @@ def load_water_system(data_dir: str | os.PathLike) -> Atoms: return atoms -def load_water_molecule_indices(data_dir: str | os.PathLike) -> np.ndarray: +def _load_water_molecule_indices(data_dir: str | os.PathLike) -> np.ndarray: """Load the per-molecule atom indices used by the NPT barostat. Returns: @@ -211,8 +211,8 @@ def run_water_npt_simulation( atoms=load_water_system(data_dir), force_field=force_field, md_integrator=MDIntegrator.NPT_MC_LANGEVIN, - molecule_indices=load_water_molecule_indices(data_dir), - **get_water_md_kwargs(run_mode), + molecule_indices=_load_water_molecule_indices(data_dir), + **_get_water_md_kwargs(run_mode), ) @@ -221,7 +221,7 @@ def run_water_npt_simulation( # -------------------------------------------------------------------------------------- -def get_solvent_md_kwargs(run_mode: RunMode) -> dict[str, Any]: +def _get_solvent_md_kwargs(run_mode: RunMode) -> dict[str, Any]: """Return the solvent simulation configuration for the given run mode.""" if run_mode == RunMode.DEV: return SOLVENT_SIMULATION_CONFIG_DEV @@ -265,7 +265,7 @@ def load_solvent_system(data_dir: str | os.PathLike, system_name: str) -> Atoms: return atoms -def load_solvent_molecule_indices( +def _load_solvent_molecule_indices( data_dir: str | os.PathLike, system_name: str ) -> np.ndarray: """Load the per-molecule atom indices for a solvent system. @@ -291,7 +291,7 @@ def run_solvent_npt_simulations( state is None if the corresponding simulation failed. """ system_names = get_solvent_system_names(run_mode) - md_kwargs = get_solvent_md_kwargs(run_mode) + md_kwargs = _get_solvent_md_kwargs(run_mode) simulation_states: list[SimulationState | None] = [] for system_name in system_names: @@ -300,7 +300,7 @@ def run_solvent_npt_simulations( atoms=load_solvent_system(data_dir, system_name), force_field=force_field, md_integrator=MDIntegrator.NPT_MC_LANGEVIN, - molecule_indices=load_solvent_molecule_indices(data_dir, system_name), + molecule_indices=_load_solvent_molecule_indices(data_dir, system_name), **md_kwargs, ) simulation_states.append(simulation_state) From 338866e84c35b3ed9cc9467eb214c34e1c4602c0 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 16:29:48 +0200 Subject: [PATCH 07/12] feat: remove file divider --- src/mlipaudit/utils/molecular_liquids.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index e3cecff6..1cb43c04 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -154,11 +154,6 @@ def average_equilibrated_density(densities: np.ndarray) -> float: return float(np.mean(densities[n_frames_equilibration:])) -# -------------------------------------------------------------------------------------- -# Water data loading + simulation -# -------------------------------------------------------------------------------------- - - def _get_water_md_kwargs(run_mode: RunMode) -> dict[str, Any]: """Return the water simulation configuration for the given run mode.""" if run_mode == RunMode.DEV: @@ -216,11 +211,6 @@ def run_water_npt_simulation( ) -# -------------------------------------------------------------------------------------- -# Solvent data loading + simulation -# -------------------------------------------------------------------------------------- - - def _get_solvent_md_kwargs(run_mode: RunMode) -> dict[str, Any]: """Return the solvent simulation configuration for the given run mode.""" if run_mode == RunMode.DEV: From ed09f711821e86e62f0a9510578ccc58c8ccb5d3 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 16:31:23 +0200 Subject: [PATCH 08/12] feat: update `data_dir` docstring --- src/mlipaudit/benchmark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mlipaudit/benchmark.py b/src/mlipaudit/benchmark.py index 0afa836a..f86a7bfc 100644 --- a/src/mlipaudit/benchmark.py +++ b/src/mlipaudit/benchmark.py @@ -236,7 +236,7 @@ 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 (e.g. an RDF and a density benchmark running the same system). + input data. """ return self.data_input_dir / (self.data_name or self.name) From 1b10eaf74f10df9be13934dadd0c2ec9392c4b14 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 16:38:22 +0200 Subject: [PATCH 09/12] feat: more private functions --- src/mlipaudit/utils/molecular_liquids.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mlipaudit/utils/molecular_liquids.py b/src/mlipaudit/utils/molecular_liquids.py index 1cb43c04..2f32a7bc 100644 --- a/src/mlipaudit/utils/molecular_liquids.py +++ b/src/mlipaudit/utils/molecular_liquids.py @@ -163,7 +163,7 @@ def _get_water_md_kwargs(run_mode: RunMode) -> dict[str, Any]: return WATER_SIMULATION_CONFIG -def load_water_system(data_dir: str | os.PathLike) -> Atoms: +def _load_water_system(data_dir: str | os.PathLike) -> Atoms: """Load the water box structure from the given data directory. Returns: @@ -203,7 +203,7 @@ def run_water_npt_simulation( """ logger.info("Running water box NPT simulation.") return run_simulation( - atoms=load_water_system(data_dir), + atoms=_load_water_system(data_dir), force_field=force_field, md_integrator=MDIntegrator.NPT_MC_LANGEVIN, molecule_indices=_load_water_molecule_indices(data_dir), @@ -243,7 +243,7 @@ def get_solvent_molecule_indices_file_name(system_name: str) -> str: return f"{system_name}_molecule_indices.npy" -def load_solvent_system(data_dir: str | os.PathLike, system_name: str) -> Atoms: +def _load_solvent_system(data_dir: str | os.PathLike, system_name: str) -> Atoms: """Load a solvent structure from the given data directory. Returns: @@ -287,7 +287,7 @@ def run_solvent_npt_simulations( for system_name in system_names: logger.info("Running NPT simulation for %s.", system_name) simulation_state = run_simulation( - atoms=load_solvent_system(data_dir, system_name), + atoms=_load_solvent_system(data_dir, system_name), force_field=force_field, md_integrator=MDIntegrator.NPT_MC_LANGEVIN, molecule_indices=_load_solvent_molecule_indices(data_dir, system_name), From a8c914109847724d2442afd05647bdeae3544929 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 16:45:42 +0200 Subject: [PATCH 10/12] docs: add references --- .../benchmarks/molecular_liquids/radial_distribution.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/benchmarks/molecular_liquids/radial_distribution.rst b/docs/source/benchmarks/molecular_liquids/radial_distribution.rst index f7119f3f..6d896558 100644 --- a/docs/source/benchmarks/molecular_liquids/radial_distribution.rst +++ b/docs/source/benchmarks/molecular_liquids/radial_distribution.rst @@ -54,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 ` 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 `, 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 From 9953bca33550ce11d51cc081f90534dd95838e50 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 17:32:37 +0200 Subject: [PATCH 11/12] docs: add CHANGELOG entry for molecular liquids split; simplify solvent score mean - Add a Release 0.1.5 CHANGELOG entry describing the RDF/density split, the new water_density / solvent_density benchmarks, the shared-simulation reuse, and the new Benchmark.data_name mechanism. - Drop the unreachable `r.score if r.score is not None else 0.0` guard in the two solvent benchmarks' score aggregation; the per-structure `score` field defaults to 0.0 and is never None. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 20 +++++++++++++++++++ .../solvent_density/solvent_density.py | 4 +--- .../solvent_radial_distribution.py | 4 +--- 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 721c7406..8f95b7cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,25 @@ # 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. Each now + contributes its own score, leaderboard column, and UI page. Density is scored in + its own right (relative deviation of the equilibrium density from the experimental + reference) rather than only being displayed on the radial-distribution pages. +- Run the shared NPT simulation only once per system group: the radial-distribution + and density benchmarks of a group share a `reusable_output_id` and identical + `ModelOutput` signatures, so the density benchmark reuses the trajectory instead of + repeating the simulation when both are run together. +- 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; + the density benchmarks reuse the radial-distribution input data without a separate + upload. +- 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 diff --git a/src/mlipaudit/benchmarks/solvent_density/solvent_density.py b/src/mlipaudit/benchmarks/solvent_density/solvent_density.py index bc88e6b5..a7ca08cb 100644 --- a/src/mlipaudit/benchmarks/solvent_density/solvent_density.py +++ b/src/mlipaudit/benchmarks/solvent_density/solvent_density.py @@ -215,7 +215,5 @@ def analyze(self) -> SolventDensityResult: for structure in structure_results if structure.density_deviation is not None ), - score=statistics.mean( - r.score if r.score is not None else 0.0 for r in structure_results - ), + score=statistics.mean(r.score for r in structure_results), ) diff --git a/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py b/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py index 5572264f..b92cbcc5 100644 --- a/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py +++ b/src/mlipaudit/benchmarks/solvent_radial_distribution/solvent_radial_distribution.py @@ -268,9 +268,7 @@ def analyze(self) -> SolventRadialDistributionResult: for structure in structure_results if structure.peak_deviation is not None ), - score=statistics.mean( - r.score if r.score is not None else 0.0 for r in structure_results - ), + score=statistics.mean(r.score for r in structure_results), ) @property From 95a8be7880defcb0c41cd53a133c7dc1dc4284f7 Mon Sep 17 00:00:00 2001 From: lwalew Date: Wed, 8 Jul 2026 17:34:32 +0200 Subject: [PATCH 12/12] docs: shorten release --- CHANGELOG.md | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f95b7cb..0f00fc3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,18 +4,9 @@ - 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. Each now - contributes its own score, leaderboard column, and UI page. Density is scored in - its own right (relative deviation of the equilibrium density from the experimental - reference) rather than only being displayed on the radial-distribution pages. -- Run the shared NPT simulation only once per system group: the radial-distribution - and density benchmarks of a group share a `reusable_output_id` and identical - `ModelOutput` signatures, so the density benchmark reuses the trajectory instead of - repeating the simulation when both are run together. + `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; - the density benchmarks reuse the radial-distribution input data without a separate - upload. + 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.