Feat: add DeepMD pretrain model (#12)

* feat: add deepmd pretrain model

* separate model into different file in externals

* chore: add deepmd dependency

* skip hf model download for testing external fork

* change class name

* add HF HTTPError

* chore: downgrade deepmd to match pt version

* chore: try install deepmd from repo

* skip missing json on leaderboard; add installation instruction

* fix callout render

* fix readme path in pyproject.toml

---------

Co-authored-by: Yuan Chiang <cyrusyc@berkeley.edu>

.github/README.md

@@ -12,6 +12,25 @@
 
 MLIP Arena is a platform for evaluating foundation machine learning interatomic potentials (MLIPs) beyond conventional energy and force error metrics. It focuses on revealing the underlying physics and chemistry learned by these models and assessing their performance in molecular dynamics (MD) simulations. The platform's benchmarks are specifically designed to evaluate the readiness and reliability of open-source, open-weight models in accurately reproducing both qualitative and quantitative behaviors of atomic systems.
 
+## Installation
+
+### From PyPI (without model running capability)
+
+```bash
+pip install mlip-arena
+```
+
+### From source
+
+```bash
+git clone https://github.com/atomind-ai/mlip-arena.git
+pip install torch==2.2.0
+bash scripts/install-pyg.sh
+bash scripts/install-dgl.sh
+pip install .[test]
+pip install .[mace]
+```
+
 ## Contribute
 
 MLIP Arena is now in pre-alpha. If you're interested in joining the effort, please reach out to Yuan at [cyrusyc@berkeley.edu](mailto:cyrusyc@berkeley.edu). See [project page](https://github.com/orgs/atomind-ai/projects/1) for some outstanding tasks. 
@@ -22,18 +41,25 @@ MLIP Arena is now in pre-alpha. If you're interested in joining the effort, plea
 streamlit run serve/app.py
 ```
 
-### Add new benchmark tasks
+### Add new benchmark tasks (WIP)
+
+> [!NOTE]
+> Please reuse or extend the general tasks defined as Prefect / Atomate2 workflow. 
+> The following are some tasks implemented:
+> - [Prefect structure optimization (OPT)](../mlip_arena/tasks/optimize.py)
+> - [Prefect molecular dynamics (MD)](../mlip_arena/tasks/md.py)
+> - [Prefect equation of states (EOS)](../mlip_arena/tasks/eos/run.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.
 
-### Add new MLIP models
+### Add new MLIP models 
 
 If you have pretrained MLIP models that you would like to contribute to the MLIP Arena and show benchmark in real-time, there are two ways:
 
 #### External ASE Calculator (easy)
 
-1. Implement new ASE Calculator class in [mlip_arena/models/external.py](../mlip_arena/models/externals.py). 
+1. Implement new ASE Calculator class in [mlip_arena/models/externals](../mlip_arena/models/externals). 
 2. Name your class with awesome model name and add the same name to [registry](../mlip_arena/models/registry.yaml) with metadata.
 
 > [!CAUTION] 

.github/workflows/test.yaml

@@ -28,14 +28,14 @@ jobs:
         pip install torch==2.2.0
         bash scripts/install-pyg.sh
         bash scripts/install-dgl.sh
-        pip install .[mace]
         pip install .[test]
-        pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
+        pip install .[mace]
 
     - name: List dependencies
       run: pip list
       
     - name: Login huggingface
+      if: ${{ github.event.pull_request.head.repo.full_name == github.repository }}
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN_READ_ONLY }}
       run:

mlip_arena/models/externals/deepmd.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+import requests
+from deepmd.calculator import DP as DPCalculator
+
+from mlip_arena.models.utils import get_freer_device
+
+with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
+    REGISTRY = yaml.safe_load(f)
+
+class DeepMD(DPCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["DeepMD"]["checkpoint"],
+        device=None,
+        **kwargs,
+    ):
+        device = device or get_freer_device()
+
+        cache_dir = Path.home() / ".cache" / "deepmd"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        model_path = cache_dir / checkpoint
+
+        url = "https://bohrium-api.dp.tech/ds-dl/mlip-arena-tfpk-v1.zip"
+
+        if not model_path.exists():
+            import zipfile
+
+            print(f"Downloading DeepMD model from {url} to {model_path}...")
+            try:
+                response = requests.get(url, stream=True, timeout=120)
+                response.raise_for_status()
+                with open(cache_dir/"temp.zip", "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                print("Download completed.")
+                with zipfile.ZipFile(cache_dir/"temp.zip", "r") as zip_ref:
+                    zip_ref.extractall(cache_dir)
+                print("Unzip completed.")
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError("Failed to download DeepMD model.") from e
+
+        
+        super().__init__(model_path, device=device, **kwargs)

mlip_arena/models/registry.yaml

@@ -245,4 +245,22 @@ ALIGNN:
   npt: true
   github: https://github.com/usnistgov/alignn
   doi: https://doi.org/10.1038/s41524-021-00650-1
-  date: 2021-11-15
+  date: 2021-11-15
+
+DeepMD:
+  module: externals
+  class: DeepMD
+  family: deepmd
+  package: deepmd-kit==v3.0.0b4
+  checkpoint: dp0808c_v024mixu.pth
+  username: 
+  last-update: 2024-10-09T00:00:00
+  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
+  datasets:
+    - MPTrj # TODO: fake HF dataset repo
+  github: https://github.com/deepmodeling/deepmd-kit/
+  doi: https://arxiv.org/abs/2312.15492
+  date: 2024-10-09
+  prediction: EFS
+  nvt: true
+  npt: true

pyproject.toml

@@ -3,13 +3,13 @@ requires=["flit_core >=3.2,<4"]
 build-backend="flit_core.buildapi"
 
 [project]
-name="mlip_arena"
+name="mlip-arena"
 version="0.0.1a1"
 authors=[
     {name="Yuan Chiang", email="cyrusyc@lbl.gov"},
 ]
-description=""
-readme="README.md"
+description="Fair and transparent benchmark of machine-learned interatomic potentials (MLIPs), beyond basic error metrics"
+readme=".github/README.md"
 requires-python=">=3.10"
 keywords=[
     "pytorch",
@@ -66,9 +66,11 @@ test = [
     "fairchem-core==1.2.0",
     "sevenn==0.9.3.post1",
     "orb-models==0.3.1",
+    "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168",
     "alignn==2024.5.27",
     "pytest",
-    "prefect>=3.0.4"
+    "prefect>=3.0.4",
+    "deepmd-kit@git+https://github.com/deepmodeling/deepmd-kit.git@v3.0.0b4"
 ]
 mace = [
     "mace-torch==0.3.4",

serve/leaderboard.py

@@ -7,21 +7,21 @@ import streamlit as st
 from mlip_arena.models import REGISTRY as MODELS
 from mlip_arena.tasks import REGISTRY as TASKS
 
+# Read the data
 DATA_DIR = Path("mlip_arena/tasks/diatomics")
 
-dfs = [
-    pd.read_json(DATA_DIR / MODELS[model].get("family") / "homonuclear-diatomics.json")
-    for model in MODELS
-]
+dfs = []
+for model in MODELS:
+    fpath = DATA_DIR / MODELS[model].get("family") / "homonuclear-diatomics.json"
+    if fpath.exists():
+        dfs.append(pd.read_json(fpath))
 df = pd.concat(dfs, ignore_index=True)
 
-
+# Create a table
 table = pd.DataFrame(
     columns=[
         "Model",
         "Element Coverage",
-        # "No. of reversed forces",
-        # "Energy-consistent forces",
         "Prediction",
         "NVT",
         "NPT",
@@ -39,8 +39,6 @@ for model in MODELS:
     new_row = {
         "Model": model,
         "Element Coverage": len(rows["name"].unique()),
-        # "No. of reversed forces": None,  # Replace with actual logic if available
-        # "Energy-consistent forces": None,  # Replace with actual logic if available
         "Prediction": metadata.get("prediction", None),
         "NVT": "✅" if metadata.get("nvt", False) else "❌",
         "NPT": "✅" if metadata.get("npt", False) else "❌",
@@ -122,10 +120,12 @@ for task in TASKS:
         # 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.")
-    
+        st.write(
+            "Rank metrics are not available yet but the task has been implemented. Please see the following task page for more information."
+        )
+
     st.page_link(
         f"tasks/{TASKS[task]['task-page']}.py",
         label="Task page",
         icon=":material/link:",
-    )
+    )

tests/test_external_calculators.py

@@ -3,6 +3,8 @@ from ase import Atoms
 
 from mlip_arena.models import MLIPEnum
 
+from requests import HTTPError
+from huggingface_hub.errors import LocalTokenNotFoundError
 
 @pytest.mark.parametrize("model", MLIPEnum)
 def test_calculate(model: MLIPEnum):
@@ -10,7 +12,12 @@ def test_calculate(model: MLIPEnum):
     if model.name == "ALIGNN":
         pytest.xfail("ALIGNN has poor file download mechanism")
 
-    calc = MLIPEnum[model.name].value()
+    try:
+        calc = MLIPEnum[model.name].value()
+
+    except (LocalTokenNotFoundError, HTTPError):
+        # Gracefully skip the test if HF_TOKEN is not available
+        pytest.skip("Skipping test because HF_TOKEN is not available for downloading the model.")
 
     atoms = Atoms(
         "OO",