Major cleanup (#51)

* clean up; add mattersim combustion; add mace-mpa

* add app test

* only push to hf after test passes on main

* add streamlit to test deps

* add test badge

* update readme

.github/README.md

@@ -1,7 +1,8 @@
 <div align="center">
     <h1>MLIP Arena</h1>
+    <a href="https://github.com/atomind-ai/mlip-arena/actions"><img alt="GitHub Actions Workflow Status" src="https://img.shields.io/github/actions/workflow/status/atomind-ai/mlip-arena/test.yaml"></a>
     <a href="https://zenodo.org/doi/10.5281/zenodo.13704399"><img src="https://zenodo.org/badge/776930320.svg" alt="DOI"></a>
-    <a href="https://huggingface.co/spaces/atomind/mlip-arena"><img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo-with-title.svg" style="height: 20px; background-color: white;" alt="Hugging Face"></a>
+    <a href="https://huggingface.co/spaces/atomind/mlip-arena"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Space-blue" alt="Hugging Face"></a>
     <!-- <a href="https://discord.gg/W8WvdQtT8T"><img alt="Discord" src="https://img.shields.io/discord/1299613474820984832?logo=discord"> -->
 </a>
 </div>
@@ -107,8 +108,8 @@ streamlit run serve/app.py
 > - [Prefect molecular dynamics (MD)](../mlip_arena/tasks/md.py)
 > - [Prefect equation of states (EOS)](../mlip_arena/tasks/eos.py)
 
-1. Follow the task template to implement the task class and upload the script along with metadata to the MLIP Arena [here](../mlip_arena/tasks/README.md).
-2. Code a benchmark script to evaluate the performance of your model on the task. The script should be able to load the model and the dataset, and output the evaluation metrics.
+<!-- 1. Follow the task template to implement the task class and upload the script along with metadata to the MLIP Arena [here](../mlip_arena/tasks/README.md).
+2. Code a benchmark script to evaluate the performance of your model on the task. The script should be able to load the model and the dataset, and output the evaluation metrics. -->
 
 ### Add new MLIP models 
 
@@ -129,12 +130,10 @@ If you have pretrained MLIP models that you would like to contribute to the MLIP
 2. Follow the template to code the I/O interface for your model [here](../mlip_arena/models/README.md). 
 3. Update model [registry](../mlip_arena/models/registry.yaml) with metadata
 
-> [!NOTE] 
-> CPU benchmarking will be performed automatically. Due to the limited amount GPU compute, if you would like to be considered for GPU benchmarking, please create a pull request to demonstrate the offline performance of your model (published paper or preprint). We will review and select the models to be benchmarked on GPU.
+<!-- > [!NOTE] 
+> CPU benchmarking will be performed automatically. Due to the limited amount GPU compute, if you would like to be considered for GPU benchmarking, please create a pull request to demonstrate the offline performance of your model (published paper or preprint). We will review and select the models to be benchmarked on GPU. -->
 
-
-
-### Add new datasets
+<!-- ### Add new datasets
 
 The "ultimate" goal is to compile the copies of all the open data in a unified format for lifelong learning with [Hugging Face Auto-Train](https://huggingface.co/docs/hub/webhooks-guide-auto-retrain). 
 
@@ -150,4 +149,4 @@ The "ultimate" goal is to compile the copies of all the open data in a unified f
 #### Molecular dynamics calculations
 
 - [ ] [MD17](http://www.sgdml.org/#datasets)
-- [ ] [MD22](http://www.sgdml.org/#datasets)
+- [ ] [MD22](http://www.sgdml.org/#datasets) -->

.github/workflows/sync-hf.yaml

@@ -3,6 +3,7 @@ name: Sync to Hugging Face hub
 on:
   workflow_run:
     workflows: [Python Test]
+    branches: [main]
     types: [completed]
   workflow_dispatch:
 

.github/workflows/test.yaml

@@ -61,4 +61,4 @@ jobs:
         PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}
         PREFECT_API_URL: ${{ secrets.PREFECT_API_URL }}
       run: |
-        pytest --dist=loadscope -vra tests -n 5
+        pytest -vra -n 5 tests 

.gitignore

@@ -1,11 +1,11 @@
 *.out
-*.ipynb
 *.extxyz
 *.traj
 mlip_arena/tasks/*/
 examples/
 lab/
 manuscripts/
+datasets/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

mlip_arena/models/externals/mace-mp.py

@@ -37,3 +37,33 @@ class MACE_MP_Medium(MACECalculator):
         super().__init__(
             model_paths=model, device=device, default_dtype=default_dtype, **kwargs
         )
+
+class MACE_MPA(MACECalculator):
+    def __init__(
+        self,
+        checkpoint="https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model",
+        device: str | None = None,
+        default_dtype="float32",
+        **kwargs,
+    ):
+        cache_dir = Path.home() / ".cache" / "mace"
+        checkpoint_url_name = "".join(
+            c for c in os.path.basename(checkpoint) if c.isalnum() or c in "_"
+        )
+        cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
+        if not os.path.isfile(cached_model_path):
+            import urllib
+
+            os.makedirs(cache_dir, exist_ok=True)
+            _, http_msg = urllib.request.urlretrieve(checkpoint, cached_model_path)
+            if "Content-Type: text/html" in http_msg:
+                raise RuntimeError(
+                    f"Model download failed, please check the URL {checkpoint}"
+                )
+        model = cached_model_path
+
+        device = device or str(get_freer_device())
+
+        super().__init__(
+            model_paths=model, device=device, default_dtype=default_dtype, **kwargs
+        )

mlip_arena/models/registry.yaml

@@ -2,7 +2,7 @@ MACE-MP(M):
   module: externals
   class: MACE_MP_Medium
   family: mace-mp
-  package: mace-torch==0.3.4
+  package: mace-torch==0.3.9
   checkpoint: 2023-12-03-mace-128-L1_epoch-199.model
   username: cyrusyc
   last-update: 2024-03-25T14:30:00
@@ -60,6 +60,7 @@ M3GNet:
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
+    - stability
   github: https://github.com/materialsvirtuallab/matgl
   doi: https://doi.org/10.1038/s43588-022-00349-3
   date: 2022-02-05
@@ -85,6 +86,7 @@ MatterSim:
   gpu-tasks:
     - homonuclear-diatomics
     - stability
+    - combustion
   github: https://github.com/microsoft/mattersim
   doi: https://arxiv.org/abs/2405.04967
   date: 2024-12-05
@@ -167,6 +169,28 @@ eqV2(OMat):
   doi: https://arxiv.org/abs/2410.12771
   license: Modified Apache-2.0 (Meta)
 
+MACE-MPA:
+  module: externals
+  class: MACE_MPA
+  family: mace-mp
+  package: mace-torch==0.3.9
+  checkpoint: mace-mpa-0-medium.model
+  username: 
+  last-update: 2025-11-19T00:00:00
+  datetime: 2024-12-09T00:00:00 # TODO: Fake datetime
+  datasets:
+    - MPTrj # TODO: fake HF dataset repo
+    - Alexandria
+  gpu-tasks:
+    - homonuclear-diatomics
+    - stability
+  github: https://github.com/ACEsuit/mace
+  doi: 
+  date: 2024-12-09
+  prediction: EFS
+  nvt: true
+  npt: true
+  license: MIT
 
 EquiformerV2(OC22):
   module: externals
@@ -237,7 +261,7 @@ MACE-OFF(M):
   module: externals
   class: MACE_OFF_Medium
   family: mace-off
-  package: mace-torch==0.3.4
+  package: mace-torch==0.3.9
   checkpoint: MACE-OFF23_medium.model
   username: cyrusyc
   last-update: 2024-03-25T14:30:00
@@ -272,7 +296,7 @@ ANI2x:
   date: 2024-05-23
   prediction: EFS
   nvt: true
-  npt: false
+  npt: true
   license: MIT
 
 ALIGNN:

mlip_arena/tasks/README.md

@@ -1,8 +1,8 @@
-## Note on task registration
+<!-- ## Note on task registration
 
 1. Use `ast` to parse task classes from the uploaded script.
 2. Add the classes and their supported tasks to the task registry file `registry.yaml`.
 3. Run tests on HF Space to ensure the task is working as expected.
 4. [Push task script to the Space](https://huggingface.co/docs/huggingface_hub/guides/upload) and sync with github repository.
 5. Create task folder in [mlip-arena](https://huggingface.co/datasets/atomind/mlip-arena) HF Dataset.
-6. 
+6.  -->

mlip_arena/tasks/combustion/mattersim/hydrogen.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00c0c38af5321151ff4a3fc64935df168689030ba31cad0be2589379360b333b
+size 226556

mlip_arena/tasks/combustion/water.ipynb

@@ -4,30 +4,23 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "No module named 'deepmd'\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from pathlib import Path\n",
     "\n",
-    "from ase import units, Atoms\n",
-    "from ase.build import molecule\n",
-    "from ase.io import read, write\n",
     "from dask.distributed import Client\n",
     "from dask_jobqueue import SLURMCluster\n",
+    "from mlip_arena.models import REGISTRY, MLIPEnum\n",
+    "from mlip_arena.tasks.md import run as MD\n",
+    "from mlip_arena.tasks.utils import get_calculator\n",
     "from prefect import flow\n",
     "from prefect_dask import DaskTaskRunner\n",
-    "from pymatgen.core import Molecule\n",
-    "from pymatgen.io.packmol import PackmolBoxGen\n",
     "\n",
-    "from mlip_arena.models import REGISTRY, MLIPEnum\n",
-    "from mlip_arena.tasks.md import run as MD"
+    "from ase import Atoms, units\n",
+    "from ase.build import molecule\n",
+    "from ase.io import read, write\n",
+    "from pymatgen.core import Molecule\n",
+    "from pymatgen.io.packmol import PackmolBoxGen"
    ]
   },
   {
@@ -42,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -57,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -68,7 +61,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,15 +81,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "tolerance = 2.0\n",
     "input_gen = PackmolBoxGen(\n",
@@ -132,19 +117,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "tags": []
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "atoms = read(\"H256O128.extxyz\")\n",
     "print(atoms)"
@@ -154,40 +131,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "#!/bin/bash\n",
-      "\n",
-      "#SBATCH -A matgen\n",
-      "#SBATCH --mem=0\n",
-      "#SBATCH -t 02:00:00\n",
-      "#SBATCH -J combustion-water\n",
-      "#SBATCH -q regular\n",
-      "#SBATCH -N 1\n",
-      "#SBATCH -C gpu\n",
-      "#SBATCH -G 4\n",
-      "#SBATCH --exclusive\n",
-      "source ~/.bashrc\n",
-      "module load python\n",
-      "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.15:38781 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:187: UserWarning: Port 8787 is already in use.\n",
-      "Perhaps you already have a cluster running?\n",
-      "Hosting the HTTP server on port 44831 instead\n",
-      "  warnings.warn(\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "nodes_per_alloc = 1\n",
     "gpus_per_alloc = 4\n",
@@ -197,8 +141,8 @@
     "    cores=1,\n",
     "    memory=\"64 GB\",\n",
     "    shebang=\"#!/bin/bash\",\n",
-    "    account=\"matgen\",\n",
-    "    walltime=\"02:00:00\",\n",
+    "    account=\"m4282\",\n",
+    "    walltime=\"00:30:00\",\n",
     "    job_mem=\"0\",\n",
     "    job_script_prologue=[\n",
     "        \"source ~/.bashrc\",\n",
@@ -208,7 +152,7 @@
     "    job_directives_skip=[\"-n\", \"--cpus-per-task\", \"-J\"],\n",
     "    job_extra_directives=[\n",
     "        \"-J combustion-water\",\n",
-    "        \"-q regular\",\n",
+    "        \"-q debug\",\n",
     "        f\"-N {nodes_per_alloc}\",\n",
     "        \"-C gpu\",\n",
     "        f\"-G {gpus_per_alloc}\",\n",
@@ -221,7 +165,7 @@
     "\n",
     "\n",
     "print(cluster.job_script())\n",
-    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
+    "cluster.adapt(minimum_jobs=1, maximum_jobs=1)\n",
     "client = Client(cluster)"
    ]
   },
@@ -236,18 +180,23 @@
     "    futures = []\n",
     "\n",
     "    for model in MLIPEnum:\n",
+    "        if model.name != \"MatterSim\":\n",
+    "            continue\n",
+    "\n",
     "        future = MD.submit(\n",
     "            atoms=atoms,\n",
-    "            calculator_name=model,\n",
-    "            calculator_kwargs=None,\n",
+    "            calculator=get_calculator(\n",
+    "                calculator_name=model,\n",
+    "                calculator_kwargs=None,\n",
+    "            ),\n",
     "            ensemble=\"nvt\",\n",
     "            dynamics=\"nose-hoover\",\n",
     "            time_step=None,\n",
-    "            ase_md_kwargs=dict(ttime=25 * units.fs, pfactor=None),\n",
+    "            dynamics_kwargs=dict(ttime=25 * units.fs, pfactor=None),\n",
     "            total_time=1000_000,\n",
     "            temperature=[300, 3000, 3000, 300],\n",
     "            pressure=None,\n",
-    "            md_velocity_seed=0,\n",
+    "            velocity_seed=0,\n",
     "            traj_file=Path(REGISTRY[model.name][\"family\"])\n",
     "            / f\"{model.name}_{atoms.get_chemical_formula()}.traj\",\n",
     "            traj_interval=1000,\n",
@@ -269,6 +218,54 @@
    "source": [
     "results = combustion(atoms)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def combustion(atoms: Atoms):\n",
+    "    futures = []\n",
+    "\n",
+    "    for model in MLIPEnum:\n",
+    "        if model.name != \"MatterSim\":\n",
+    "            continue\n",
+    "\n",
+    "        future = MD(\n",
+    "            atoms=atoms,\n",
+    "            calculator=get_calculator(\n",
+    "                calculator_name=model,\n",
+    "                calculator_kwargs=None,\n",
+    "            ),\n",
+    "            ensemble=\"nvt\",\n",
+    "            dynamics=\"nose-hoover\",\n",
+    "            time_step=None,\n",
+    "            dynamics_kwargs=dict(ttime=25 * units.fs, pfactor=None),\n",
+    "            total_time=1000_000,\n",
+    "            temperature=[300, 3000, 3000, 300],\n",
+    "            pressure=None,\n",
+    "            velocity_seed=0,\n",
+    "            traj_file=Path(REGISTRY[model.name][\"family\"])\n",
+    "            / f\"{model.name}_{atoms.get_chemical_formula()}.traj\",\n",
+    "            traj_interval=1000,\n",
+    "            restart=True,\n",
+    "        )\n",
+    "\n",
+    "        futures.append(future)\n",
+    "\n",
+    "    return [future.result() for future in futures]\n",
+    "\n",
+    "\n",
+    "results = combustion(atoms)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

mlip_arena/tasks/diatomics/mace-mp/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1f6f4a2d4f36071625db988dde933674cdf5478951cf227e6eacc0d818c13a1f
-size 1915573
+oid sha256:9ad34875760232ee25a34b7d6e8a54d75e3b8ddd38efaf5de3aa7a3d6e19474d
+size 3837066

mlip_arena/tasks/diatomics/run.ipynb

@@ -14,17 +14,15 @@
     "\n",
     "import numpy as np\n",
     "import pandas as pd\n",
-    "from ase import Atom, Atoms\n",
-    "from ase.data import chemical_symbols, covalent_radii, vdw_alvarez\n",
-    "from ase.io import read, write\n",
-    "from pymatgen.core import Element\n",
     "from scipy import stats\n",
-    "from scipy.interpolate import splrep, BSpline, CubicSpline, UnivariateSpline\n",
+    "from scipy.interpolate import UnivariateSpline\n",
     "from tqdm.auto import tqdm\n",
     "\n",
-    "from mlip_arena.models import MLIPEnum, REGISTRY\n",
-    "\n",
-    "%matplotlib inline"
+    "from ase import Atom, Atoms\n",
+    "from ase.data import chemical_symbols, covalent_radii, vdw_alvarez\n",
+    "from ase.io import read, write\n",
+    "from mlip_arena.models import REGISTRY, MLIPEnum\n",
+    "from pymatgen.core import Element"
    ]
   },
   {
@@ -48,18 +46,16 @@
    "outputs": [],
    "source": [
     "for model in MLIPEnum:\n",
-    "    \n",
     "    model_name = model.name\n",
-    "    \n",
-    "    if model_name != 'MatterSim':\n",
+    "\n",
+    "    if model_name != \"MACE-MPA\":\n",
     "        continue\n",
-    "    \n",
+    "\n",
     "    print(f\"========== {model_name} ==========\")\n",
     "\n",
     "    calc = MLIPEnum[model_name].value()\n",
     "\n",
     "    for symbol in tqdm(chemical_symbols[1:]):\n",
-    "\n",
     "        s = set([symbol])\n",
     "\n",
     "        if \"X\" in s:\n",
@@ -68,13 +64,14 @@
     "        try:\n",
     "            atom = Atom(symbol)\n",
     "            rmin = 0.9 * covalent_radii[atom.number]\n",
-    "            rvdw = vdw_alvarez.vdw_radii[atom.number] if atom.number < len(vdw_alvarez.vdw_radii) else np.nan\n",
+    "            rvdw = (\n",
+    "                vdw_alvarez.vdw_radii[atom.number]\n",
+    "                if atom.number < len(vdw_alvarez.vdw_radii)\n",
+    "                else np.nan\n",
+    "            )\n",
     "            rmax = 3.1 * rvdw if not np.isnan(rvdw) else 6\n",
     "            rstep = 0.01\n",
-    "\n",
-    "            a = 2 * rmax\n",
-    "\n",
-    "            npts = int((rmax - rmin)/rstep)\n",
+    "            npts = int((rmax - rmin) / rstep)\n",
     "\n",
     "            rs = np.linspace(rmin, rmax, npts)\n",
     "            es = np.zeros_like(rs)\n",
@@ -92,15 +89,15 @@
     "                m = element.valence[1]\n",
     "                if element.valence == (0, 2):\n",
     "                    m = 0\n",
-    "            except:\n",
+    "            except Exception:\n",
     "                m = 0\n",
     "\n",
-    "\n",
+    "            a = 2 * rmax\n",
     "            r = rs[0]\n",
     "\n",
     "            positions = [\n",
-    "                [a/2-r/2, a/2, a/2],\n",
-    "                [a/2+r/2, a/2, a/2],\n",
+    "                [a / 2 - r / 2, a / 2, a / 2],\n",
+    "                [a / 2 + r / 2, a / 2, a / 2],\n",
     "            ]\n",
     "\n",
     "            traj_fpath = out_dir / f\"{model_name}.extxyz\"\n",
@@ -115,8 +112,8 @@
     "                    da,\n",
     "                    positions=positions,\n",
     "                    # magmoms=magmoms,\n",
-    "                    cell=[a, a+0.001, a+0.002],\n",
-    "                    pbc=True\n",
+    "                    cell=[a, a + 0.001, a + 0.002],\n",
+    "                    pbc=True,\n",
     "                )\n",
     "\n",
     "            print(atoms)\n",
@@ -124,13 +121,12 @@
     "            atoms.calc = calc\n",
     "\n",
     "            for i, r in enumerate(tqdm(rs)):\n",
-    "\n",
     "                if i < skip:\n",
     "                    continue\n",
     "\n",
     "                positions = [\n",
-    "                    [a/2-r/2, a/2, a/2],\n",
-    "                    [a/2+r/2, a/2, a/2],\n",
+    "                    [a / 2 - r / 2, a / 2, a / 2],\n",
+    "                    [a / 2 + r / 2, a / 2, a / 2],\n",
     "                ]\n",
     "\n",
     "                # atoms.set_initial_magnetic_moments(magmoms)\n",
@@ -162,48 +158,47 @@
    },
    "outputs": [],
    "source": [
-    "\n",
-    "\n",
     "for model in MLIPEnum:\n",
-    "    \n",
     "    model_name = model.name\n",
-    "    \n",
+    "\n",
     "    # if model_name != \"MatterSim\":\n",
     "    #     continue\n",
     "\n",
     "    print(f\"========== {model_name} ==========\")\n",
-    "    \n",
-    "    df = pd.DataFrame(columns=[\n",
-    "        \"name\", \n",
-    "        \"method\", \n",
-    "        \"R\", \"E\", \"F\", \"S^2\",\n",
-    "        \"force-flip-times\",\n",
-    "        \"force-total-variation\",\n",
-    "        \"force-jump\",\n",
-    "        \"energy-diff-flip-times\",\n",
-    "        \"energy-grad-norm-max\",\n",
-    "        \"energy-jump\",\n",
-    "        \"energy-total-variation\",\n",
-    "        \"tortuosity\",\n",
-    "        \"conservation-deviation\",\n",
-    "        \"spearman-descending-force\",\n",
-    "        \"spearman-ascending-force\",\n",
-    "        \"spearman-repulsion-energy\",\n",
-    "        \"spearman-attraction-energy\",\n",
-    "        \"pbe-energy-mae\",\n",
-    "        \"pbe-force-mae\"\n",
-    "    ])\n",
-    "    \n",
     "\n",
-    "    for symbol in tqdm(chemical_symbols[1:]):\n",
+    "    df = pd.DataFrame(\n",
+    "        columns=[\n",
+    "            \"name\",\n",
+    "            \"method\",\n",
+    "            \"R\",\n",
+    "            \"E\",\n",
+    "            \"F\",\n",
+    "            \"S^2\",\n",
+    "            \"force-flip-times\",\n",
+    "            \"force-total-variation\",\n",
+    "            \"force-jump\",\n",
+    "            \"energy-diff-flip-times\",\n",
+    "            \"energy-grad-norm-max\",\n",
+    "            \"energy-jump\",\n",
+    "            \"energy-total-variation\",\n",
+    "            \"tortuosity\",\n",
+    "            \"conservation-deviation\",\n",
+    "            \"spearman-descending-force\",\n",
+    "            \"spearman-ascending-force\",\n",
+    "            \"spearman-repulsion-energy\",\n",
+    "            \"spearman-attraction-energy\",\n",
+    "            \"pbe-energy-mae\",\n",
+    "            \"pbe-force-mae\",\n",
+    "        ]\n",
+    "    )\n",
     "\n",
+    "    for symbol in tqdm(chemical_symbols[1:]):\n",
     "        da = symbol + symbol\n",
     "\n",
     "        out_dir = Path(REGISTRY[model_name][\"family\"]) / da\n",
     "\n",
     "        traj_fpath = out_dir / f\"{model_name}.extxyz\"\n",
     "\n",
-    "\n",
     "        if traj_fpath.exists():\n",
     "            traj = read(traj_fpath, index=\":\")\n",
     "        else:\n",
@@ -211,11 +206,10 @@
     "\n",
     "        Rs, Es, Fs, S2s = [], [], [], []\n",
     "        for atoms in traj:\n",
-    "\n",
     "            vec = atoms.positions[1] - atoms.positions[0]\n",
     "            r = np.linalg.norm(vec)\n",
     "            e = atoms.get_potential_energy()\n",
-    "            f = np.inner(vec/r, atoms.get_forces()[1])\n",
+    "            f = np.inner(vec / r, atoms.get_forces()[1])\n",
     "            # s2 = np.mean(np.power(atoms.get_magnetic_moments(), 2))\n",
     "\n",
     "            Rs.append(r)\n",
@@ -243,33 +237,36 @@
     "\n",
     "        # avoid numerical sensitity close to zero\n",
     "        rounded_fs = np.copy(fs)\n",
-    "        rounded_fs[np.abs(rounded_fs) < 1e-2] = 0 # 10meV/A\n",
+    "        rounded_fs[np.abs(rounded_fs) < 1e-2] = 0  # 10meV/A\n",
     "        fs_sign = np.sign(rounded_fs)\n",
     "        mask = fs_sign != 0\n",
     "        rounded_fs = rounded_fs[mask]\n",
     "        fs_sign = fs_sign[mask]\n",
     "        f_flip = np.diff(fs_sign) != 0\n",
-    "        \n",
+    "\n",
     "        fdiff = np.diff(fs)\n",
     "        fdiff_sign = np.sign(fdiff)\n",
     "        mask = fdiff_sign != 0\n",
     "        fdiff = fdiff[mask]\n",
     "        fdiff_sign = fdiff_sign[mask]\n",
     "        fdiff_flip = np.diff(fdiff_sign) != 0\n",
-    "        fjump = np.abs(fdiff[:-1][fdiff_flip]).sum() + np.abs(fdiff[1:][fdiff_flip]).sum()\n",
-    "        \n",
+    "        fjump = (\n",
+    "            np.abs(fdiff[:-1][fdiff_flip]).sum() + np.abs(fdiff[1:][fdiff_flip]).sum()\n",
+    "        )\n",
     "\n",
     "        ediff = np.diff(es)\n",
-    "        ediff[np.abs(ediff) < 1e-3] = 0 # 1meV\n",
+    "        ediff[np.abs(ediff) < 1e-3] = 0  # 1meV\n",
     "        ediff_sign = np.sign(ediff)\n",
     "        mask = ediff_sign != 0\n",
     "        ediff = ediff[mask]\n",
     "        ediff_sign = ediff_sign[mask]\n",
     "        ediff_flip = np.diff(ediff_sign) != 0\n",
-    "        ejump = np.abs(ediff[:-1][ediff_flip]).sum() + np.abs(ediff[1:][ediff_flip]).sum()\n",
-    "        \n",
+    "        ejump = (\n",
+    "            np.abs(ediff[:-1][ediff_flip]).sum() + np.abs(ediff[1:][ediff_flip]).sum()\n",
+    "        )\n",
+    "\n",
     "        try:\n",
-    "            pbe_traj = read(f'./vasp/{da}/PBE.extxyz', index=\":\")\n",
+    "            pbe_traj = read(f\"./vasp/{da}/PBE.extxyz\", index=\":\")\n",
     "\n",
     "            pbe_rs, pbe_es, pbe_fs = [], [], []\n",
     "\n",
@@ -278,7 +275,7 @@
     "                r = np.linalg.norm(vec)\n",
     "                pbe_rs.append(r)\n",
     "                pbe_es.append(atoms.get_potential_energy())\n",
-    "                pbe_fs.append(np.inner(vec/r, atoms.get_forces()[1]))\n",
+    "                pbe_fs.append(np.inner(vec / r, atoms.get_forces()[1]))\n",
     "\n",
     "            pbe_rs = np.array(pbe_rs)\n",
     "            pbe_es = np.array(pbe_es)\n",
@@ -302,43 +299,9 @@
     "            print(e)\n",
     "            pbe_energy_mae = None\n",
     "            pbe_force_mae = None\n",
-    "        \n",
-    "        \n",
-    "#         edged_es = np.convolve(es, [1, -2, 1], mode='valid')\n",
-    "#         # edged_es[np.abs(edged_es) < 0.1] = 0\n",
-    "#         prob = np.exp(-es[1:-1]) / np.sum(np.exp(-es[1:-1]))\n",
-    "#         edged_es *= prob\n",
-    "#         # edged_es /= np.abs(es[1:-1])\n",
-    "#         ejump = np.linalg.norm(edged_es)\n",
-    "#         ejump = np.abs(edged_es).sum() / 2.0\n",
-    "        \n",
-    "#         edged_fs = np.convolve(fs, [1, -2, 1], mode='valid')\n",
-    "#         # edged_fs[np.abs(edged_fs) < 0.1] = 0\n",
-    "#         edged_fs *= prob\n",
-    "#         fjump = np.linalg.norm(edged_fs)\n",
-    "        # fjump = np.abs(edged_fs).sum() / 2.0\n",
-    "        \n",
-    "#         fig, axes = plt.subplot_mosaic(\n",
-    "#             \"\"\"\n",
-    "#             ac\n",
-    "#             bd\n",
-    "#             \"\"\",\n",
-    "#             constrained_layout=True\n",
-    "#         )\n",
-    "        \n",
-    "\n",
-    "#         axes['a'].plot(rs, es)\n",
-    "#         axes['b'].plot(rs[1:-1], edged_es)\n",
-    "#         # axes['b'].plot(0.5*(rs[1:] + rs[:-1]), np.diff(es))\n",
-    "#         axes['b'].text(0.7, 0.7, f\"{ejump:.3e}\", transform=axes['b'].transAxes)\n",
-    "        \n",
-    "#         axes['c'].plot(rs, fs)\n",
-    "#         axes['d'].plot(rs[1:-1], edged_fs)\n",
-    "#         axes['d'].text(0.7, 0.7, f\"{fjump:.3e}\", transform=axes['d'].transAxes)\n",
-    "        \n",
     "\n",
     "        conservation_deviation = np.mean(np.abs(fs + de_dr))\n",
-    "        \n",
+    "\n",
     "        etv = np.sum(np.abs(np.diff(es)))\n",
     "\n",
     "        data = {\n",
@@ -358,12 +321,20 @@
     "            \"energy-total-variation\": etv,\n",
     "            \"tortuosity\": etv / (abs(es[0] - es.min()) + (es[-1] - es.min())),\n",
     "            \"conservation-deviation\": conservation_deviation,\n",
-    "            \"spearman-descending-force\": stats.spearmanr(rs[iminf:], fs[iminf:]).statistic,\n",
-    "            \"spearman-ascending-force\": stats.spearmanr(rs[:iminf], fs[:iminf]).statistic,\n",
-    "            \"spearman-repulsion-energy\": stats.spearmanr(rs[imine:], es[imine:]).statistic,\n",
-    "            \"spearman-attraction-energy\": stats.spearmanr(rs[:imine], es[:imine]).statistic,\n",
+    "            \"spearman-descending-force\": stats.spearmanr(\n",
+    "                rs[iminf:], fs[iminf:]\n",
+    "            ).statistic,\n",
+    "            \"spearman-ascending-force\": stats.spearmanr(\n",
+    "                rs[:iminf], fs[:iminf]\n",
+    "            ).statistic,\n",
+    "            \"spearman-repulsion-energy\": stats.spearmanr(\n",
+    "                rs[imine:], es[imine:]\n",
+    "            ).statistic,\n",
+    "            \"spearman-attraction-energy\": stats.spearmanr(\n",
+    "                rs[:imine], es[:imine]\n",
+    "            ).statistic,\n",
     "            \"pbe-energy-mae\": pbe_energy_mae,\n",
-    "            \"pbe-force-mae\": pbe_force_mae\n",
+    "            \"pbe-force-mae\": pbe_force_mae,\n",
     "        }\n",
     "\n",
     "        df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
@@ -373,37 +344,17 @@
     "    if json_fpath.exists():\n",
     "        df0 = pd.read_json(json_fpath)\n",
     "        df = pd.concat([df0, df], ignore_index=True)\n",
-    "        df.drop_duplicates(inplace=True, subset=[\"name\", \"method\"], keep='last')\n",
+    "        df.drop_duplicates(inplace=True, subset=[\"name\", \"method\"], keep=\"last\")\n",
     "\n",
     "    df.to_json(json_fpath, orient=\"records\")"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e0dd4367-3dca-440f-a7a9-7fdd84183f2c",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "df"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4e6ae884-89f3-43f2-8fd9-19bf00c91566",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "mlip-arena",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "mlip-arena"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

mlip_arena/tasks/md.py

@@ -143,7 +143,7 @@ def _get_ensemble_schedule(
         isinstance(pressure, np.ndarray) and pressure.ndim == 1
     ):
         p_schedule = _interpolate_quantity(pressure, n_steps)
-    elif isinstance(pressure, np.ndarray) and pressure.ndim == 4:
+    elif isinstance(pressure, np.ndarray) and pressure.ndim == 3:
         p_schedule = interp1d(np.arange(n_steps + 1), pressure, kind="linear")
         assert isinstance(p_schedule, np.ndarray)
     else:

pyproject.toml

@@ -63,7 +63,8 @@ test = [
     "pytest",
     "pytest-xdist",
     "prefect==3.1.11",
-    "pymatgen>=2025.1.9"
+    "pymatgen>=2025.1.9",
+    "streamlit==1.38.0"
 ]
 mace = [
     "mace-torch==0.3.9",

serve/leaderboard.py

@@ -119,16 +119,15 @@ for task in TASKS:
 
     #  Call the function from the imported module
     if hasattr(task_module, "render"):
+        st.page_link(
+            f"tasks/{TASKS[task]['task-page']}.py",
+            label="Go to the associated task page",
+            icon=":material/link:",
+        )
         task_module.render()
         # if st.button(f"Go to task page"):
         #     st.switch_page(f"tasks/{TASKS[task]['task-page']}.py")
     else:
         st.write(
-            "Rank metrics are not available yet but the task has been implemented. Please see the following task page for more information."
+            "Rank metrics are not available yet but the task has been implemented. Please see the task page for more information."
         )
-
-    st.page_link(
-        f"tasks/{TASKS[task]['task-page']}.py",
-        label="Go to the associated task page",
-        icon=":material/link:",
-    )

serve/ranks/homonuclear-diatomics.py

@@ -173,3 +173,4 @@ def render():
             - **Force flips**: The number of force direction changes.
             """
         )
+        st.info('PBE energies and forces are provided __only__ for reference. Due to the known convergence issue of plane-wave DFT with diatomic molecules and different dataset the models might be trained on, comparing models with PBE is not rigorous and thus these metrics are excluded from rank aggregation.', icon=":material/warning:")

serve/tasks/combustion.py

@@ -6,7 +6,6 @@ import plotly.colors as pcolors
 import plotly.express as px
 import plotly.graph_objects as go
 import streamlit as st
-
 from mlip_arena.models import REGISTRY as MODELS
 
 DATA_DIR = Path("mlip_arena/tasks/combustion")
@@ -36,6 +35,7 @@ models = container.multiselect(
         "ORBv2",
         "EquiformerV2(OC20)",
         "eSCN(OC20)",
+        "MatterSim",
     ],
 )
 
@@ -64,7 +64,9 @@ if not models:
 def get_data(models):
     # List comprehension for concise looping and filtering
     dfs = [
-        pd.read_json(DATA_DIR / MODELS[str(model)]["family"].lower() / "hydrogen.json")[lambda df: df["method"] == model]
+        pd.read_json(DATA_DIR / MODELS[str(model)]["family"].lower() / "hydrogen.json")[
+            lambda df: df["method"] == model
+        ]
         for model in models
     ]
     # Concatenate all filtered DataFrames
@@ -177,8 +179,8 @@ st.plotly_chart(fig)
 
 # Energy
 
-exp_ref = -68.3078 # kcal/mol
-factor = 23.0609 
+exp_ref = -68.3078  # kcal/mol
+factor = 23.0609
 nh2os = 128
 
 fig = go.Figure()
@@ -205,10 +207,12 @@ target_steps = df["target_steps"].iloc[0]
 fig.add_shape(
     go.layout.Shape(
         type="line",
-        x0=0, x1=target_steps,
-        y0=exp_ref, y1=exp_ref,  # y-values for the horizontal line
+        x0=0,
+        x1=target_steps,
+        y0=exp_ref,
+        y1=exp_ref,  # y-values for the horizontal line
         line=dict(color="Red", width=2, dash="dash"),
-        layer="below"
+        layer="below",
     )
 )
 
@@ -281,28 +285,36 @@ st.plotly_chart(fig)
 fig = go.Figure()
 
 
-df["reaction_energy"] = df["energies"].apply(lambda x: x[-1] - x[0]) / nh2os * factor # kcal/mol
+df["reaction_energy"] = (
+    df["energies"].apply(lambda x: x[-1] - x[0]) / nh2os * factor
+)  # kcal/mol
 
 df["reaction_energy_abs_err"] = np.abs(df["reaction_energy"] - exp_ref)
 
 df.sort_values("reaction_energy_abs_err", inplace=True)
 
-fig.add_traces([
-    go.Bar(
-        x=df["method"],
-        y=df["reaction_energy"],
-        marker=dict(color=[method_color_mapping[method] for method in df["method"]]),
-        text=[f"{y:.2f}" for y in df["reaction_energy"]],
-    ),
-])
+fig.add_traces(
+    [
+        go.Bar(
+            x=df["method"],
+            y=df["reaction_energy"],
+            marker=dict(
+                color=[method_color_mapping[method] for method in df["method"]]
+            ),
+            text=[f"{y:.2f}" for y in df["reaction_energy"]],
+        ),
+    ]
+)
 
 fig.add_shape(
     go.layout.Shape(
         type="line",
-        x0=-0.5, x1=len(df["method"]) - 0.5,  # range covering the bars
-        y0=exp_ref, y1=exp_ref,  # y-values for the horizontal line
+        x0=-0.5,
+        x1=len(df["method"]) - 0.5,  # range covering the bars
+        y0=exp_ref,
+        y1=exp_ref,  # y-values for the horizontal line
         line=dict(color="Red", width=2, dash="dash"),
-        layer="below"
+        layer="below",
     )
 )
 
@@ -356,7 +368,7 @@ fig.add_trace(
 fig.update_layout(
     title="Reaction yield (2H2 + O2 -> 2H2O, 64 units)",
     xaxis_title="Yield (%)",
-    yaxis_title="Method"
+    yaxis_title="Method",
 )
 
 st.plotly_chart(fig)
@@ -433,7 +445,6 @@ for method in df_exploded["method"].unique():
             ),
             marker=dict(color=method_color_mapping[method], size=3),
             showlegend=True,
-            
         ),
     )
 
@@ -564,5 +575,4 @@ st.markdown("""
 [1] Hasche, A., Navid, A., Krause, H., & Eckart, S. (2023). Experimental and numerical assessment of the effects of hydrogen admixtures on premixed methane-oxygen flames. Fuel, 352, 128964.
             
 [2] Lide, D. R. (Ed.). (2004). CRC handbook of chemistry and physics (Vol. 85). CRC press.
-"""
-)
+""")

serve/tasks/homonuclear-diatomics.py

@@ -5,11 +5,10 @@ import pandas as pd
 import plotly.colors as pcolors
 import plotly.graph_objects as go
 import streamlit as st
-from ase.data import chemical_symbols
+from mlip_arena.models import REGISTRY
 from plotly.subplots import make_subplots
-from scipy.interpolate import CubicSpline
 
-from mlip_arena.models import REGISTRY
+from ase.data import chemical_symbols
 
 st.markdown(
     """
@@ -30,10 +29,24 @@ valid_models = [
 mlip_methods = container.multiselect(
     "MLIPs",
     valid_models,
-    ["MACE-MP(M)", "CHGNet", "M3GNet", "MatterSim", "SevenNet", "ORBv2", "eqV2(OMat)", "ANI2x"],
+    [
+        "MACE-MP(M)",
+        "CHGNet",
+        "M3GNet",
+        "MatterSim",
+        "SevenNet",
+        "ORBv2",
+        "eqV2(OMat)",
+        "ANI2x",
+    ],
 )
 dft_methods = container.multiselect("DFT Methods", ["PBE"], ["PBE"])
 
+container.info(
+    "PBE energies and forces are provided __only__ for reference. Due to the known convergence issue of plane-wave DFT with diatomic molecules and different dataset the models might be trained on, comparing models with PBE is not rigorous and thus these metrics are excluded from rank aggregation.",
+    icon=":material/warning:",
+)
+
 st.markdown("### Settings")
 vis = st.container(border=True)
 energy_plot = vis.checkbox("Show energy curves", value=True)
@@ -119,11 +132,10 @@ def get_plots(df, energy_plot: bool, force_plot: bool, method_color_mapping: dic
             rs = rs[ind]
             es = es[ind]
             fs = fs[ind]
-            
+
             # if method not in ["PBE"]:
             es = es - es[-1]
 
-
             # if method in ["PBE"]:
             #     xs = np.linspace(rs.min() * 0.99, rs.max() * 1.01, int(5e2))
             # else:

tests/test_app.py

@@ -0,0 +1,27 @@
+import streamlit as st
+from streamlit.testing.v1 import AppTest
+import pytest
+from pathlib import Path
+
+path = Path(__file__).parents[1] / "serve" 
+
+@pytest.fixture
+def home():
+    at = AppTest.from_file(str(path / "app.py"), default_timeout=60)
+    at.run()
+    assert not at.exception
+    return at
+
+def test_leaderboard(home):
+    # Test the leaderboard page by simulating navigation.
+    at = home.switch_page(str(path / "leaderboard.py"))
+    assert not at.exception
+    
+def test_task_pages(home):
+    # Test each task page using the TASKS registry.
+    from mlip_arena.tasks import REGISTRY as TASKS
+
+    for task, details in TASKS.items():
+        page_path = str(path / f"tasks/{details['task-page']}.py")
+        at = home.switch_page(page_path)
+        assert not at.exception