Fix MD kwargs (#52)

* add pypi badge

* move task page buttom

* fix md kwargs

* patch stability input module

* add one script linux installation

* add list of implemented tasks

* enforce release using workflow dispatch

.github/README.md

@@ -1,6 +1,7 @@
 <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://pypi.org/project/mlip-arena/"><img alt="PyPI - Version" src="https://img.shields.io/pypi/v/mlip-arena"></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://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"> -->
@@ -30,6 +31,15 @@ pip install mlip-arena
 **Linux**
 
 ```bash
+# (Optional) Install uv
+curl -LsSf https://astral.sh/uv/install.sh | sh
+source $HOME/.local/bin/env
+# One script uv pip installation
+bash scripts/install-linux.sh
+```
+
+```bash
+# Or from command line
 git clone https://github.com/atomind-ai/mlip-arena.git
 cd mlip-arena
 pip install torch==2.2.0
@@ -87,6 +97,19 @@ for model in MLIPEnum:
     results.append(result)
 ```
 
+### List of implemented tasks
+
+The implemented tasks are available under `mlip_arena.tasks.<module>.run` or `from mlip_arena.tasks import *` for convenient imports (currently doesn't work if [phonopy](https://phonopy.github.io/phonopy/install.html) is not installed).
+
+- [OPT](../mlip_arena/tasks/optimize.py#L56): Structure optimization
+- [EOS](../mlip_arena/tasks/eos.py#L42): Equation of state (energy-volume scan)
+- [MD](../mlip_arena/tasks/md.py#L200): Molecular dynamics with flexible dynamics (NVE, NVT, NPT) and temperature/pressure scheduling (annealing, shearing, *etc*)
+- [PHONON](../mlip_arena/tasks/phonon.py#L110): Phonon calculation driven by [phonopy](https://phonopy.github.io/phonopy/install.html)
+- [NEB](../mlip_arena/tasks/neb.py#L96): Nudged elastic band
+- [NEB_FROM_ENDPOINTS](../mlip_arena/tasks/neb.py#L164): Nudge elastic band with convenient image interpolation (linear or IDPP)
+- [ELASTICITY](../mlip_arena/tasks/elasticity.py#L78): Elastic tensor calculation
+
+
 ## 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). 

.github/workflows/release.yaml

@@ -1,74 +1,20 @@
 name: Publish Release
 
 on:
-  workflow_run:
-    workflows: [Python Test]
-    types: [completed]
+  # workflow_run:
+  #   workflows: [Python Test]
+  #   branches: [main]
+  #   types: [completed]
   workflow_dispatch:
 
 permissions:
   contents: write  # Ensure write access to push tags
 
 jobs:
-  release:
-    name: Create GitHub Release
-    runs-on: ubuntu-latest
-
-    steps:
-    # Step 1: Checkout the code
-    - name: Checkout code
-      uses: actions/checkout@v3
-
-    # Step 2: Set up Python
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.x'
-
-    # Step 3: Install dependencies
-    - name: Install dependencies
-      run: pip install toml
-
-    # Step 4: Extract version from pyproject.toml
-    - name: Extract version
-      id: get_version
-      run: |
-        VERSION=$(python -c "import toml; print(toml.load('pyproject.toml')['project']['version'])")
-        echo "VERSION=$VERSION" >> $GITHUB_ENV
-
-    # Step 5: Check if tag exists on remote
-    - name: Check if tag exists on remote
-      id: check_tag
-      run: |
-        if git ls-remote --tags origin | grep "refs/tags/v${{ env.VERSION }}"; then
-          echo "Tag v${{ env.VERSION }} already exists on remote."
-          echo "tag_exists=true" >> $GITHUB_ENV
-        else
-          echo "tag_exists=false" >> $GITHUB_ENV
-        fi
-
-    # Step 6: Create and push a new tag (if it doesn't exist)
-    - name: Create Git tag
-      if: env.tag_exists == 'false'
-      run: |
-        git config --global user.name "github-actions[bot]"
-        git config --global user.email "github-actions[bot]@users.noreply.github.com"
-        git tag -a "v${{ env.VERSION }}" -m "Release v${{ env.VERSION }}"
-        git push origin "v${{ env.VERSION }}"
-
-    # Step 7: Create GitHub release (if tag didn't exist)
-    - name: Create GitHub Release
-      if: env.tag_exists == 'false'
-      uses: softprops/action-gh-release@v1
-      with:
-        tag_name: "v${{ env.VERSION }}"
-      env:
-        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  
   pypi:
     name: Publish to PyPI
     runs-on: ubuntu-latest
-    needs: release  # This job runs after the release job
+    # needs: release  # This job runs after the release job
 
     steps:
     # Step 1: Checkout the code

mlip_arena/tasks/md.py

@@ -157,29 +157,29 @@ def _get_ensemble_defaults(
     dynamics: str | MolecularDynamics,
     t_schedule: np.ndarray,
     p_schedule: np.ndarray,
-    ase_md_kwargs: dict | None = None,
+    dynamics_kwargs: dict | None = None,
 ) -> dict:
     """Update ASE MD kwargs"""
-    ase_md_kwargs = ase_md_kwargs or {}
+    dynamics_kwargs = dynamics_kwargs or {}
 
     if ensemble == "nve":
-        ase_md_kwargs.pop("temperature", None)
-        ase_md_kwargs.pop("temperature_K", None)
-        ase_md_kwargs.pop("externalstress", None)
+        dynamics_kwargs.pop("temperature", None)
+        dynamics_kwargs.pop("temperature_K", None)
+        dynamics_kwargs.pop("externalstress", None)
     elif ensemble == "nvt":
-        ase_md_kwargs["temperature_K"] = t_schedule[0]
-        ase_md_kwargs.pop("externalstress", None)
+        dynamics_kwargs["temperature_K"] = t_schedule[0]
+        dynamics_kwargs.pop("externalstress", None)
     elif ensemble == "npt":
-        ase_md_kwargs["temperature_K"] = t_schedule[0]
-        ase_md_kwargs["externalstress"] = p_schedule[0]  # * 1e3 * units.bar
+        dynamics_kwargs["temperature_K"] = t_schedule[0]
+        dynamics_kwargs["externalstress"] = p_schedule[0]  # * 1e3 * units.bar
 
     if isinstance(dynamics, str) and dynamics.lower() == "langevin":
-        ase_md_kwargs["friction"] = ase_md_kwargs.get(
+        dynamics_kwargs["friction"] = dynamics_kwargs.get(
             "friction",
             10.0 * 1e-3 / units.fs,  # Same default as in VASP: 10 ps^-1
         )
 
-    return ase_md_kwargs
+    return dynamics_kwargs
 
 
 def _generate_task_run_name():
@@ -206,8 +206,8 @@ def run(
     total_time: float = 1000,  # fs
     temperature: float | Sequence | np.ndarray | None = 300.0,  # K
     pressure: float | Sequence | np.ndarray | None = None,  # eV/A^3
-    ase_md_kwargs: dict | None = None,
-    md_velocity_seed: int | None = None,
+    dynamics_kwargs: dict | None = None,
+    velocity_seed: int | None = None,
     zero_linear_momentum: bool = True,
     zero_angular_momentum: bool = True,
     traj_file: str | Path | None = None,
@@ -235,12 +235,12 @@ def run(
         pressure=pressure,
     )
 
-    ase_md_kwargs = _get_ensemble_defaults(
+    dynamics_kwargs = _get_ensemble_defaults(
         ensemble=ensemble,
         dynamics=dynamics,
         t_schedule=t_schedule,
         p_schedule=p_schedule,
-        ase_md_kwargs=ase_md_kwargs,
+        dynamics_kwargs=dynamics_kwargs,
     )
 
     if isinstance(dynamics, str):
@@ -289,7 +289,7 @@ def run(
                     MaxwellBoltzmannDistribution(
                         atoms=atoms,
                         temperature_K=t_schedule[last_step],
-                        rng=np.random.default_rng(seed=md_velocity_seed),
+                        rng=np.random.default_rng(seed=velocity_seed),
                     )
 
                 if zero_linear_momentum:
@@ -303,7 +303,7 @@ def run(
                 MaxwellBoltzmannDistribution(
                     atoms=atoms,
                     temperature_K=t_schedule[last_step],
-                    rng=np.random.default_rng(seed=md_velocity_seed),
+                    rng=np.random.default_rng(seed=velocity_seed),
                 )
 
             if zero_linear_momentum:
@@ -314,7 +314,7 @@ def run(
     md_runner = md_class(
         atoms=atoms,
         timestep=time_step * units.fs,
-        **ase_md_kwargs,
+        **dynamics_kwargs,
     )
 
     if traj_file is not None:

mlip_arena/tasks/stability/input.py

@@ -0,0 +1,74 @@
+import os
+from pathlib import Path
+from typing import Generator, Iterable
+
+from huggingface_hub import HfApi, hf_hub_download
+from prefect import task
+
+from ase import Atoms
+from ase.db import connect
+from mlip_arena.tasks.utils import logger
+
+
+def save_to_db(
+    atoms_list: list[Atoms] | Iterable[Atoms] | Atoms,
+    db_path: Path | str,
+    upload: bool = True,
+    hf_token: str | None = os.getenv("HF_TOKEN", None),
+    repo_id: str = "atomind/mlip-arena",
+    repo_type: str = "dataset",
+    subfolder: str = Path(__file__).parent.name,
+):
+    """Save ASE Atoms objects to an ASE database and optionally upload to Hugging Face Hub."""
+
+    if upload and hf_token is None:
+        raise ValueError("HF_TOKEN is required to upload the database.")
+
+    db_path = Path(db_path)
+
+    if isinstance(atoms_list, Atoms):
+        atoms_list = [atoms_list]
+
+    with connect(db_path) as db:
+        for atoms in atoms_list:
+            if not isinstance(atoms, Atoms):
+                raise ValueError("atoms_list must contain ASE Atoms objects.")
+            db.write(atoms)
+
+    if upload:
+        api = HfApi(token=hf_token)
+        api.upload_file(
+            path_or_fileobj=db_path,
+            path_in_repo=f"{subfolder}/{db_path.name}",
+            repo_id=repo_id,
+            repo_type=repo_type,
+        )
+        logger.info(f"{db_path.name} uploaded to {repo_id}/{subfolder}")
+
+    return db_path
+
+
+@task
+def get_atoms_from_db(
+    db_path: Path | str,
+    hf_token: str | None = os.getenv("HF_TOKEN", None),
+    repo_id: str = "atomind/mlip-arena",
+    repo_type: str = "dataset",
+    subfolder: str = Path(__file__).parent.name,
+    force_download: bool = False,
+) -> Generator[Atoms, None, None]:
+    """Retrieve ASE Atoms objects from an ASE database."""
+    db_path = Path(db_path)
+    if not db_path.exists():
+        db_path = hf_hub_download(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            subfolder=subfolder,
+            # local_dir=db_path.parent,
+            filename=db_path.name,
+            token=hf_token,
+            force_download=force_download,
+        )
+    with connect(db_path) as db:
+        for row in db.select():
+            yield row.toatoms()

scripts/install-linux.sh

@@ -0,0 +1,8 @@
+TORCH=2.2
+CUDA=cu121
+uv pip install torch==${TORCH}.0
+uv pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
+uv pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.0+${CUDA}.html
+uv pip install dgl -f https://data.dgl.ai/wheels/torch-${TORCH}/${CUDA}/repo.html
+uv pip install -e .[test]
+uv pip install -e .[mace]

serve/leaderboard.py

@@ -117,13 +117,14 @@ for task in TASKS:
 
     task_module = importlib.import_module(f"ranks.{TASKS[task]['rank-page']}")
 
+    st.page_link(
+        f"tasks/{TASKS[task]['task-page']}.py",
+        label="Go to the associated task page",
+        icon=":material/link:",
+    )
+
     #  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")

tests/test_md.py

@@ -23,6 +23,7 @@ def test_nve(model: MLIPEnum):
         dynamics="velocityverlet",
         total_time=10,
         time_step=2,
+        dynamics_kwargs={},
     )
 
     assert isinstance(result["atoms"].get_potential_energy(), float)