.. _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 runs on a size-stratified set of protein chains. For each system it
measures, with warm-up and outlier trimming:

* **Model throughput** — the model forward pass (energy + forces). 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**, per backend — an end-to-end short **NVT** **MD** simulation at
  **300 K**, with per-step times derived from the cumulative step count reported at
  each logger call (discarding the initial compilation interval), reported
  as **ns/day**. ``mlip`` models are run on both the **JAX-MD** and **ASE** backends
  (the latter via ``MLIPForceFieldASECalculator``), while external ASE calculators are
  run on **ASE** only. Because ASE is shared across all models, the ASE numbers give an
  apples-to-apples MD comparison between JAX and external models, while JAX-MD shows
  the ``mlip`` best case.

The gap between the metrics and engines may reflect simulation overhead (neighbour lists, the
integrator and engine).

Dataset
-------

The dataset is a size-stratified set of protein chains taken from the
`PDB <https://www.rcsb.org/>`_. Chains were sourced from a PISCES cull list
(non-redundant at 25% sequence identity, resolution ≤ 2.0 Å, no chain breaks)
and a curated small-protein list, screened to charge-neutral sequences at pH 7:

* **2JOF** chain A — Trp-cage TC10b mini-protein (284 atoms)
* **1R0R** chain I — turkey ovomucoid third domain (OMTKY3) (748 atoms)
* **3TXS** chain A — bacteriophage 44RR small terminase gp16 (1513 atoms)
* **4QMD** chain A — human envoplakin plakin-repeat domain (3018 atoms)
* **6U1V** chain A — TcsD acyl-CoA dehydrogenase from FK506 biosynthesis (5964 atoms)

Interpretation
--------------

The benchmark produces a score in ``[0, 1]`` from the **model forward time** ``t`` of
each system via a Hill function ``1 / (1 + (t / t_ref(N))ᵏ)`` averaged over systems, so
faster models score higher. The forward time (rather than the MD step time) is scored.
Because ``t`` is wall-clock time, this score is
hardware-dependent and is only comparable across models run on the same GPU.

Each system's time is normalised by a **reference cost curve**
``t_ref(N) = overhead + per_atom · N`` for a system of ``N`` atoms.
A forward pass costs a size-independent overhead plus a marginal
per-atom cost. Normalising by ``t_ref(N)`` removes that size dependence, so every
system contributes comparably to the average and ``k`` controls how sharply
models separate. A model sitting exactly on the reference curve scores 0.5.

.. note::

   The scored forward time is **not** measured identically across model families, so
   the score is not a strictly fair comparison between ``mlip`` and external models.
   For ``mlip`` models it times the pure network forward on a pre-built graph, with
   neighbour-list construction excluded; for external ASE calculators it times a full
   recomputation that **includes** the calculator's neighbour-list build. External
   models therefore carry a systematic overhead that ``mlip`` models do not, so
   cross-family score comparisons should be treated with caution — the score is most
   meaningful when comparing models of the same family on identical hardware.