feat: add inference_speed benchmark with throughput score

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Changelog
 
+## Unreleased
+
+- Add an `inference_speed` benchmark that measures, per system size, both **model
+  throughput** (the raw forward pass — pure network for mlip models, forced ASE
+  recompute for external models) and **MD throughput** (end-to-end ns/day). It reuses
+  the `scaling` dataset and produces a (hardware-relative) speed score from the per-atom
+  model forward time, contributing to the overall model score. The GUI switches between
+  metrics on log–log axes with power-law fits and variance, and a summary table
+  including the model's graph cutoff.
+- Benchmarks can now reuse another benchmark's dataset via the `dataset_name` attribute.
+
 ## Release 0.1.4
 
 - Add an Apache-2.0 `LICENSE` file and declare the license in `pyproject.toml`.

docs/source/api_reference/general/inference_speed.rst

@@ -0,0 +1,20 @@
+.. _inference_speed_api:
+
+Inference Speed
+===============
+
+.. module:: mlipaudit.benchmarks.inference_speed.inference_speed
+
+.. autoclass:: InferenceSpeedBenchmark
+
+    .. automethod:: __init__
+
+    .. automethod:: run_model
+
+    .. automethod:: analyze
+
+.. autoclass:: InferenceSpeedResult
+
+.. autoclass:: InferenceSpeedStructureResult
+
+.. autoclass:: InferenceSpeedModelOutput

docs/source/api_reference/index.rst

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

docs/source/benchmarks/general/index.rst

@@ -11,3 +11,4 @@ applicable across molecular systems.
 
     Stability <stability>
     Scaling <scaling>
+    Inference Speed <inference_speed>

docs/source/benchmarks/general/inference_speed.rst

@@ -0,0 +1,46 @@
+.. _inference_speed:
+
+Inference Speed
+===============
+
+Purpose
+-------
+
+This benchmark measures how fast a machine-learned interatomic potential (**MLIP**) runs
+and how that speed scales with system size. It reports two complementary metrics so
+models can be compared on performance as well as accuracy.
+
+Description
+-----------
+
+The benchmark reuses the ``scaling`` dataset (a size-stratified set of protein chains).
+For each system it measures, with warm-up and outlier trimming:
+
+* **Model throughput** — the raw model forward pass (energy + forces), independent of
+  the simulation engine. For mlip models this is the pure network forward on a
+  pre-built graph (mirroring mlip-jax's ``scripts/time_inference.py``); for external
+  ASE calculators it is a forced recomputation on the pre-built atoms (which also
+  includes the calculator's neighbour-list construction). Reported as **atoms/s**.
+* **MD throughput** — an end-to-end short **NVT** **MD** simulation at **300 K**,
+  timed per episode (discarding the first to ignore compilation). Reported as
+  **ns/day**.
+
+The gap between the two reflects simulation overhead (neighbour lists, the integrator
+and engine). The GUI lets you switch metrics, plots them on log–log axes with
+power-law fits and per-system variance, and shows a per-model summary including the
+model's **graph cutoff** (which drives neighbour count and therefore speed).
+
+Dataset
+-------
+
+This benchmark reuses the ``scaling`` dataset — see :ref:`scaling` for details of the
+size-stratified protein chains.
+
+Interpretation
+--------------
+
+The benchmark produces a score in ``[0, 1]`` based on the per-atom **model forward
+time** ``t`` via a Hill function ``1 / (1 + (t / t₀)ᵏ)`` averaged over systems, so
+faster models score higher. The forward time (rather than the MD step time) is scored
+because it is engine-independent. Because ``t`` is wall-clock time, this score is
+hardware-dependent and is only comparable across models run on the same GPU.

src/mlipaudit/benchmark.py

@@ -104,6 +104,11 @@ class Benchmark(ABC):
 
     reusable_output_id: tuple[str, ...] | None = None
 
+    #: Name of the dataset (subdirectory and HuggingFace archive) this benchmark reads
+    #: its inputs from. Defaults to ``name``; set it to another benchmark's name to
+    #: reuse that benchmark's dataset instead of shipping a duplicate.
+    dataset_name: str | None = None
+
     def __init__(
         self,
         force_field: ForceField | ASECalculator,
@@ -224,17 +229,24 @@ def check_can_run_model(cls, force_field: ForceField) -> bool:
 
         return True
 
+    @property
+    def _dataset_name(self) -> str:
+        """The dataset directory/archive name, defaulting to the benchmark name."""
+        return self.dataset_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()
+        already_exists = (self.data_input_dir / self._dataset_name).exists()
         if not already_exists:
             hf_hub_download(
                 repo_id="InstaDeepAI/MLIPAudit-data",
-                filename=f"{self.name}.zip",
+                filename=f"{self._dataset_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"{self._dataset_name}.zip", "r"
+            ) as z:
                 z.extractall(self.data_input_dir)
 
     @abstractmethod

src/mlipaudit/benchmarks/__init__.py

@@ -34,6 +34,12 @@
     FoldingStabilityModelOutput,
     FoldingStabilityResult,
 )
+from mlipaudit.benchmarks.inference_speed.inference_speed import (
+    InferenceSpeedBenchmark,
+    InferenceSpeedModelOutput,
+    InferenceSpeedResult,
+    InferenceSpeedStructureResult,
+)
 from mlipaudit.benchmarks.noncovalent_interactions.noncovalent_interactions import (
     NoncovalentInteractionsBenchmark,
     NoncovalentInteractionsModelOutput,