Clean sync from main branch - 2025-07-05 06:02:08

.devcontainer/devcontainer.json

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

.gitattributes

@@ -0,0 +1,12 @@
+*.json filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.db filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/SevenNet.pkl filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/input.pkl filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/M3GNet.pkl filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/MACE-MPA.pkl filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/MACE-MP(M).pkl filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/MatterSim.pkl filter=lfs diff=lfs merge=lfs -text
+examples/mof/classification/ORBv2.pkl filter=lfs diff=lfs merge=lfs -text
+*.pdf filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.github/README.md

@@ -0,0 +1,198 @@
+<div align="center">
+    <h1>⚔️ MLIP Arena ⚔️</h1>
+    <a href="https://openreview.net/forum?id=ysKfIavYQE#discussion"><img alt="Static Badge" src="https://img.shields.io/badge/ICLR-AI4Mat-blue"></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://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://pypi.org/project/mlip-arena/"><img alt="PyPI - Downloads" src="https://img.shields.io/pypi/dm/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://discord.gg/W8WvdQtT8T"><img alt="Discord" src="https://img.shields.io/discord/1299613474820984832?logo=discord"> -->
+</a>
+</div>
+
+Foundation machine learning interatomic potentials (MLIPs), trained on extensive databases containing millions of density functional theory (DFT) calculations, have revolutionized molecular and materials modeling, but existing benchmarks suffer from data leakage, limited transferability, and an over-reliance on error-based metrics tied to specific density functional theory (DFT) references.
+
+We introduce MLIP Arena, a unified benchmark platform for evaluating foundation MLIP performance beyond conventional error metrics. It focuses on revealing the physical soundness learned by MLIPs and assessing their utilitarian performance agnostic to underlying model architecture and training dataset. 
+
+***By moving beyond static DFT references and revealing the important failure modes*** of current foundation MLIPs in real-world settings, MLIP Arena provides a reproducible framework to guide the next-generation MLIP development toward improved predictive accuracy and runtime efficiency while maintaining physical consistency.
+
+MLIP Arena leverages modern pythonic workflow orchestrator 💙
+ [Prefect](https://www.prefect.io/) 💙
+ to enable advanced task/flow chaining and caching.
+
+![Thumnail](../serve/assets/workflow.png)
+
+> [!NOTE]
+> Contributions of new tasks through PRs are very welcome! If you're interested in joining the effort, please reach out to Yuan at [[email protected]](mailto:[email protected]). See [project page](https://github.com/orgs/atomind-ai/projects/1) for some outstanding tasks, or propose new feature requests in [Discussion](https://github.com/atomind-ai/mlip-arena/discussions/new?category=ideas).
+
+## Announcement
+
+- **[April 8, 2025]** [🎉 **MLIP Arena is accepted as an ICLR AI4Mat Spotlight!** 🎉](https://openreview.net/forum?id=ysKfIavYQE#discussion) Huge thanks to all co-authors for their contributions!
+
+
+## Installation
+
+### From PyPI (prefect workflow only, without pretrained models)
+
+```bash
+pip install mlip-arena
+```
+
+### From source (with integrated pretrained models, advanced)
+
+> [!CAUTION] 
+> We strongly recommend clean build in a new virtual environment due to the compatibility issues between multiple popular MLIPs. We provide a single installation script using `uv` for minimal package conflicts and fast installation!
+
+> [!CAUTION]
+> To automatically download farichem OMat24 checkpoint, please make sure you have gained downloading access to their HuggingFace [***model repo***](https://huggingface.co/facebook/OMAT24) (not dataset repo), and login locally on your machine through `huggginface-cli login` (see [HF hub authentication](https://huggingface.co/docs/huggingface_hub/en/quick-start#authentication))
+
+**Linux**
+
+```bash
+# (Optional) Install uv, way faster than pip, why not? :)
+curl -LsSf https://astral.sh/uv/install.sh | sh
+source $HOME/.local/bin/env
+
+git clone https://github.com/atomind-ai/mlip-arena.git
+cd mlip-arena
+
+# One script uv pip installation
+bash scripts/install.sh
+```
+
+> [!TIP]
+> Sometimes installing all compiled models takes all the available local storage. Optional pip flag `--no-cache` could be uesed. `uv cache clean` will be helpful too.
+
+**Mac**
+
+```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-macosx.sh
+```
+
+## Quickstart
+
+### The first example: Molecular Dynamics
+
+Arena provides a unified interface to run all the compiled MLIPs. This can be achieved simply by looping through `MLIPEnum`:
+
+```python
+from mlip_arena.models import MLIPEnum
+from mlip_arena.tasks import MD 
+from mlip_arena.tasks.utils import get_calculator
+
+from ase import units
+from ase.build import bulk
+
+atoms = bulk("Cu", "fcc", a=3.6) * (5, 5, 5)
+
+results = []
+
+for model in MLIPEnum:
+    result = MD(
+        atoms=atoms,
+        calculator=get_calculator(
+            model,
+            calculator_kwargs=dict(), # passing into calculator
+            dispersion=True,
+            dispersion_kwargs=dict(
+                damping='bj', xc='pbe', cutoff=40.0 * units.Bohr
+            ), # passing into TorchDFTD3Calculator
+        ), # compatible with custom ASE Calculator
+        ensemble="nve", # nvt, nvt available
+        dynamics="velocityverlet", # compatible with any ASE Dynamics objects and their class names
+        total_time=1e3, # 1 ps = 1e3 fs
+        time_step=2, # fs
+    )
+    results.append(result)
+```
+
+### 🚀 Parallelize Benchmarks at Scale
+
+To run multiple benchmarks in parallel, add `.submit` before the task function and wrap all the tasks into a flow to dispatch the tasks to worker for concurrent execution. See Prefect Doc on [tasks](https://docs.prefect.io/v3/develop/write-tasks) and [flow](https://docs.prefect.io/v3/develop/write-flows) for more details.
+
+```python
+...
+from prefect import flow
+
+@flow
+def run_all_tasks:
+
+    futures = []
+    for model in MLIPEnum:
+        future = MD.submit(
+            atoms=atoms,
+            ...
+        )
+        future.append(future)
+
+    return [f.result(raise_on_failure=False) for f in futures]
+```
+
+For a more practical example using HPC resources, please now refer to [MD stability benchmark](../benchmarks/stability/temperature.ipynb).
+
+### 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 and Development
+
+PRs are welcome. Please clone the repo and submit PRs with changes.
+
+To make change to huggingface space, fetch large files from git lfs first and run streamlit:
+
+```
+git lfs fetch --all
+git lfs pull
+streamlit run serve/app.py
+```
+
+### Add new benchmark tasks (WIP)
+
+> [!NOTE]
+> Please reuse, extend, or chain the general tasks defined [above](#list-of-implemented-tasks)
+
+### 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/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] 
+> Remove unneccessary outputs under `results` class attributes to avoid error for MD simulations. Please refer to [CHGNet](../mlip_arena/models/externals/chgnet.py) as an example.
+
+#### Hugging Face Model (recommended, difficult)
+
+0. Inherit Hugging Face [ModelHubMixin](https://huggingface.co/docs/huggingface_hub/en/package_reference/mixins) class to your awesome model class definition. We recommend [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/en/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin).
+1. Create a new [Hugging Face Model](https://huggingface.co/new) repository and upload the model file using [push_to_hub function](https://huggingface.co/docs/huggingface_hub/en/package_reference/mixins#huggingface_hub.ModelHubMixin.push_to_hub).
+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
+
+## Citation
+
+If you find the work useful, please consider citing the following:
+
+```bibtex
+@inproceedings{
+    chiang2025mlip,
+    title={{MLIP} Arena: Advancing Fairness and Transparency in Machine Learning Interatomic Potentials through an Open and Accessible Benchmark Platform},
+    author={Yuan Chiang and Tobias Kreiman and Elizabeth Weaver and Ishan Amin and Matthew Kuner and Christine Zhang and Aaron Kaplan and Daryl Chrzan and Samuel M Blau and Aditi S. Krishnapriyan and Mark Asta},
+    booktitle={AI for Accelerated Materials Design - ICLR 2025},
+    year={2025},
+    url={https://openreview.net/forum?id=ysKfIavYQE}
+}
+```

.github/workflows/release.yaml

@@ -0,0 +1,96 @@
+name: Publish Release
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: write  # Ensure write access to push tags
+
+jobs:
+  pypi:
+    name: Publish to PyPI
+    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 requests
+
+    # Step 4: Extract current version from pyproject.toml
+    - name: Extract current version
+      id: get_version
+      run: |
+        VERSION=$(python -c "import toml; print(toml.load('pyproject.toml')['project']['version'])")
+        echo "VERSION=$VERSION" >> $GITHUB_ENV
+
+    # Step 5: Get latest version from PyPI
+    - name: Get latest version from PyPI
+      id: get_pypi_version
+      run: |
+        LATEST_PYPI_VERSION=$(python -c "import toml; import requests; PACKAGE_NAME = toml.load('pyproject.toml')['project']['name']; response = requests.get(f'https://pypi.org/pypi/{PACKAGE_NAME}/json'); print(response.json()['info']['version'])")
+        echo "LATEST_PYPI_VERSION=$LATEST_PYPI_VERSION" >> $GITHUB_ENV
+
+    # Step 6: Compare current version with the latest tag
+    - name: Check if version is bumped
+      id: check_version
+      run: |
+        if [ "${{ env.VERSION }}" = "${{ env.LATEST_PYPI_VERSION }}" ]; then
+          echo "Version not bumped. Exiting."
+          echo "version_bumped=false" >> $GITHUB_ENV
+        else
+          echo "Version bumped. Proceeding."
+          echo "version_bumped=true" >> $GITHUB_ENV
+        fi
+    
+    # Step 5: Remove problematic optional dependencies
+    - name: Strip problematic optional dependencies
+      run: |
+        python - <<EOF
+        import toml
+        from pathlib import Path
+
+        pyproject_path = Path("pyproject.toml")
+        data = toml.loads(pyproject_path.read_text())
+        
+        # Process optional dependencies
+        optional_deps = data.get("project", {}).get("optional-dependencies", {})
+        for key, deps in optional_deps.items():
+            new_deps = []
+            for dep in deps:
+                if "@git" in dep:
+                    dep = dep.split("@git")[0].strip()  # Remove everything after "@git"
+                new_deps.append(dep)
+            optional_deps[key] = new_deps
+            
+        pyproject_path.write_text(toml.dumps(data))
+        EOF
+
+    # Step 7: Install Flit (only if version bumped)
+    - name: Install Flit
+      if: env.version_bumped == 'true'
+      run: pip install flit
+
+    # Step 8: Create .pypirc file (only if version bumped)
+    - name: Create .pypirc file
+      if: env.version_bumped == 'true'
+      run: |
+        echo "[pypi]" > ~/.pypirc
+        echo "username = __token__" >> ~/.pypirc
+        echo "password = ${{ secrets.PYPI_API_TOKEN }}" >> ~/.pypirc
+
+    # Step 9: Build and publish package (only if version bumped)
+    - name: Build and Publish Package
+      if: env.version_bumped == 'true'
+      run: |
+        flit build
+        flit publish

.github/workflows/sync-hf.yaml

@@ -0,0 +1,39 @@
+name: Sync to Hugging Face hub
+
+on:
+  workflow_run:
+    workflows: [Python Test]
+    branches: [main]
+    types: [completed]
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    if: ${{ github.event.workflow_run.conclusion == 'success' }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+          
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          # Configure Git user identity
+          git config user.name "github-actions[ci]"
+          git config user.email "github-actions[ci]@users.noreply.github.com"
+
+          # Configure LFS tracking
+          git lfs track "*.pdf"
+          git lfs track "*.png"
+
+          # Create a new orphan branch (no history)
+          git checkout --orphan hf-clean
+
+          git add .
+          git commit -m "Clean sync from main branch - $(date '+%Y-%m-%d %H:%M:%S')"
+
+          # Force push to Hugging Face main branch
+          git push -f https://HF_USERNAME:[email protected]/spaces/atomind/mlip-arena hf-clean:main

.github/workflows/test.yaml

@@ -0,0 +1,103 @@
+name: Python Test
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+env:
+  UV_SYSTEM_PYTHON: 1
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+    - name: Checkout PR with full history
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v6
+      with:
+        enable-cache: true
+        cache-dependency-glob: "pyproject.toml"
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: bash scripts/install-linux.sh
+
+    - name: List dependencies
+      run: pip list
+
+    - name: Login to Hugging Face
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN_READ_ONLY }}
+      run: huggingface-cli login --token $HF_TOKEN
+
+    - name: Run tests
+      env:
+        PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}
+        PREFECT_API_URL: ${{ secrets.PREFECT_API_URL }}
+      run: pytest -vra -n 5 --dist=loadscope tests
+
+    - name: Squash commits and trial push to Hugging Face
+      if: github.event_name == 'pull_request'
+      id: trial_push
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        TRIAL_BRANCH: trial-sync-${{ github.sha }}-${{ matrix.python-version }}
+      run: |
+        # Configure Git user identity
+        git config user.name "github-actions[ci]"
+        git config user.email "github-actions[ci]@users.noreply.github.com"
+
+        # Install Git LFS
+        sudo apt-get update
+        sudo apt-get install -y git-lfs
+        git lfs install
+
+        # Configure LFS tracking for binary files (only for HF push)
+        git lfs track "*.pdf"
+        git lfs track "*.png"
+
+        git add .gitattributes
+
+        # Setup LFS for the remote
+        git lfs fetch
+        git lfs checkout
+
+        # Rebase and squash all PR commits into one
+        BASE=$(git merge-base origin/main HEAD)
+        git reset --soft $BASE
+        
+        # Re-add all files (binary files will now be tracked by LFS)
+        git add .
+        git commit -m "Squashed commit from PR #${{ github.event.pull_request.number }}"
+
+        # Create a new orphan branch (no history)
+        git checkout --orphan hf-clean
+
+        git add .
+        git commit -m "Clean sync from main branch - $(date '+%Y-%m-%d %H:%M:%S')"
+
+        # Push to temporary branch on Hugging Face
+        git push -f https://HF_USERNAME:[email protected]/spaces/atomind/mlip-arena HEAD:refs/heads/$TRIAL_BRANCH
+
+    - name: Delete trial branch from Hugging Face
+      if: steps.trial_push.outcome == 'success'
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        TRIAL_BRANCH: trial-sync-${{ github.sha }}-${{ matrix.python-version }}
+      run: |
+        git push https://HF_USERNAME:[email protected]/spaces/atomind/mlip-arena --delete $TRIAL_BRANCH || true

.gitignore

@@ -0,0 +1,169 @@
+*.out
+*.extxyz
+*.traj
+mlip_arena/tasks/*/
+benchmarks/
+lab/
+manuscripts/
+datasets/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

.streamlit/config.toml

@@ -0,0 +1,2 @@
+[server]
+fileWatcherType = "poll"

CITATION.cff

@@ -0,0 +1,23 @@
+# This CITATION.cff file was generated with cffinit.
+# Visit https://bit.ly/cffinit to generate yours today!
+
+cff-version: 1.2.0
+title: MLIP Arena
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Yuan
+    family-names: Chiang
+    email: [email protected]
+    affiliation: Lawrence Berkeley National Laboratory
+    orcid: 'https://orcid.org/0000-0002-4017-7084'
+repository-code: 'https://github.com/atomind-ai/mlip-arena'
+keywords:
+  - Quantum Chemistry
+  - Foundation Model
+  - Interatomic Potentials
+  - Machine Learning
+  - Force Fields
+license: Apache-2.0

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,14 @@
+---
+title: MLIP Arena
+emoji: ⚛
+sdk: streamlit
+sdk_version: 1.43.2 # The latest supported version
+python_version: 3.11
+app_file: serve/app.py
+colorFrom: indigo
+colorTo: yellow
+pinned: true
+short_description: Benchmark machine learning interatomic potential at scale
+---
+
+

benchmarks/c2db/ALIGNN.parquet

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

benchmarks/c2db/CHGNet.parquet

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

benchmarks/c2db/M3GNet.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32e1517a85a1b64f12fb262a0948a95be58c69edde133ce7ddf683154b8f2a95
+size 290358

benchmarks/c2db/MACE-MP(M).parquet

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

benchmarks/c2db/MACE-MPA.parquet

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

benchmarks/c2db/MatterSim.parquet

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

benchmarks/c2db/ORBv2.parquet

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

benchmarks/c2db/SevenNet.parquet

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

benchmarks/c2db/analysis.ipynb

@@ -0,0 +1,408 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0625f0a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import numpy as np\n",
+    "from ase.db import connect\n",
+    "\n",
+    "random.seed(0)\n",
+    "\n",
+    "DATA_DIR = Path(\".\")\n",
+    "\n",
+    "db = connect(DATA_DIR / \"c2db.db\")\n",
+    "random_indices = random.sample(range(1, len(db) + 1), 1000)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "005708b9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import itertools\n",
+    "\n",
+    "import pandas as pd\n",
+    "import phonopy\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "from mlip_arena.models import MLIPEnum\n",
+    "\n",
+    "for row, model in tqdm(\n",
+    "    itertools.product(db.select(filter=lambda r: r[\"id\"] in random_indices), MLIPEnum)\n",
+    "):\n",
+    "    uid = row[\"uid\"]\n",
+    "\n",
+    "    if Path(f\"{model.name}.parquet\").exists():\n",
+    "        df = pd.read_parquet(f\"{model.name}.parquet\")\n",
+    "        if uid in df[\"uid\"].unique():\n",
+    "            continue\n",
+    "    else:\n",
+    "        df = pd.DataFrame(columns=[\"model\", \"uid\", \"eigenvalues\", \"frequencies\"])\n",
+    "\n",
+    "    try:\n",
+    "        path = Path(model.name) / uid\n",
+    "        phonon = phonopy.load(path / \"phonopy.yaml\")\n",
+    "        frequencies = phonon.get_frequencies(q=(0, 0, 0))\n",
+    "\n",
+    "        data = np.load(path / \"elastic.npz\")\n",
+    "\n",
+    "        eigenvalues = data[\"eigenvalues\"]\n",
+    "\n",
+    "        new_row = pd.DataFrame(\n",
+    "            [\n",
+    "                {\n",
+    "                    \"model\": model.name,\n",
+    "                    \"uid\": uid,\n",
+    "                    \"eigenvalues\": eigenvalues,\n",
+    "                    \"frequencies\": frequencies,\n",
+    "                }\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "        df = pd.concat([df, new_row], ignore_index=True)\n",
+    "        df.drop_duplicates(subset=[\"model\", \"uid\"], keep=\"last\", inplace=True)\n",
+    "\n",
+    "        df.to_parquet(f\"{model.name}.parquet\", index=False)\n",
+    "    except Exception:\n",
+    "        pass\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b8d87638",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "uids = []\n",
+    "stabilities = []\n",
+    "for row in db.select(filter=lambda r: r[\"id\"] in random_indices):\n",
+    "    stable = row.key_value_pairs[\"dyn_stab\"]\n",
+    "    if stable.lower() == \"unknown\":\n",
+    "        stable = None\n",
+    "    else:\n",
+    "        stable = True if stable.lower() == \"yes\" else False\n",
+    "    uids.append(row.key_value_pairs[\"uid\"])\n",
+    "    stabilities.append(stable)\n",
+    "\n",
+    "\n",
+    "stabilities = np.array(stabilities)\n",
+    "\n",
+    "(stabilities == True).sum(), (stabilities == False).sum(), (stabilities == None).sum()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a3c516a7",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 104,
+   "id": "0052d0ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from matplotlib import pyplot as plt\n",
+    "from sklearn.metrics import (\n",
+    "    ConfusionMatrixDisplay,\n",
+    "    classification_report,\n",
+    "    confusion_matrix,\n",
+    ")\n",
+    "\n",
+    "from mlip_arena.models import MLIPEnum\n",
+    "\n",
+    "thres = -1e-7\n",
+    "\n",
+    "select_models = [\n",
+    "    \"ALIGNN\",\n",
+    "    \"CHGNet\",\n",
+    "    \"M3GNet\",\n",
+    "    \"MACE-MP(M)\",\n",
+    "    \"MACE-MPA\",\n",
+    "    \"MatterSim\",\n",
+    "    \"ORBv2\",\n",
+    "    \"SevenNet\",\n",
+    "]\n",
+    "\n",
+    "with plt.style.context(\"default\"):\n",
+    "    # plt.rcParams.update({\n",
+    "    #     # \"title.fontsize\": 10,\n",
+    "    #     \"axes.titlesize\": 10,\n",
+    "    #     \"axes.labelsize\": 8,\n",
+    "    # })\n",
+    "\n",
+    "    SMALL_SIZE = 8\n",
+    "    MEDIUM_SIZE = 10\n",
+    "    BIGGER_SIZE = 12\n",
+    "    plt.rcParams.update(\n",
+    "        {\n",
+    "            \"font.size\": SMALL_SIZE,\n",
+    "            \"axes.titlesize\": MEDIUM_SIZE,\n",
+    "            \"axes.labelsize\": MEDIUM_SIZE,\n",
+    "            \"xtick.labelsize\": MEDIUM_SIZE,\n",
+    "            \"ytick.labelsize\": MEDIUM_SIZE,\n",
+    "            \"legend.fontsize\": SMALL_SIZE,\n",
+    "            \"figure.titlesize\": BIGGER_SIZE,\n",
+    "        }\n",
+    "    )\n",
+    "\n",
+    "    fig, axs = plt.subplots(\n",
+    "        nrows=int(np.ceil(len(MLIPEnum) / 4)),\n",
+    "        ncols=4,\n",
+    "        figsize=(6, 3 * int(np.ceil(len(select_models) / 4))),\n",
+    "        sharey=True,\n",
+    "        sharex=True,\n",
+    "        layout=\"constrained\",\n",
+    "    )\n",
+    "    axs = axs.flatten()\n",
+    "    plot_idx = 0\n",
+    "\n",
+    "    for model in MLIPEnum:\n",
+    "        fpath = DATA_DIR / f\"{model.name}.parquet\"\n",
+    "        if not fpath.exists():\n",
+    "            continue\n",
+    "\n",
+    "        if model.name not in select_models:\n",
+    "            continue\n",
+    "\n",
+    "        df = pd.read_parquet(fpath)\n",
+    "        df[\"eigval_min\"] = df[\"eigenvalues\"].apply(\n",
+    "            lambda x: x.min() if np.isreal(x).all() else thres\n",
+    "        )\n",
+    "        df[\"freq_min\"] = df[\"frequencies\"].apply(\n",
+    "            lambda x: x.min() if np.isreal(x).all() else thres\n",
+    "        )\n",
+    "        df[\"dyn_stab\"] = ~np.logical_or(\n",
+    "            df[\"eigval_min\"] < thres, df[\"freq_min\"] < thres\n",
+    "        )\n",
+    "\n",
+    "        arg = np.argsort(uids)\n",
+    "        uids_sorted = np.array(uids)[arg]\n",
+    "        stabilities_sorted = stabilities[arg]\n",
+    "\n",
+    "        sorted_df = (\n",
+    "            df[df[\"uid\"].isin(uids_sorted)].set_index(\"uid\").reindex(uids_sorted)\n",
+    "        )\n",
+    "        mask = ~(stabilities_sorted == None)\n",
+    "\n",
+    "        y_true = stabilities_sorted[mask].astype(\"int\")\n",
+    "        y_pred = sorted_df[\"dyn_stab\"][mask].fillna(-1).astype(\"int\")\n",
+    "        cm = confusion_matrix(y_true, y_pred, labels=[1, 0, -1])\n",
+    "\n",
+    "        ax = axs[plot_idx]\n",
+    "        ConfusionMatrixDisplay(\n",
+    "            cm, display_labels=[\"stable\", \"unstable\", \"missing\"]\n",
+    "        ).plot(ax=ax, cmap=\"Blues\", colorbar=False)\n",
+    "\n",
+    "        ax.set_title(model.name)\n",
+    "        ax.set_xlabel(\"Predicted\")\n",
+    "        ax.set_ylabel(\"True\")\n",
+    "        ax.set_xticks([0, 1, 2])\n",
+    "        ax.set_xticklabels([\"stable\", \"unstable\", \"missing\"])\n",
+    "        ax.set_yticks([0, 1, 2])\n",
+    "        ax.set_yticklabels([\"stable\", \"unstable\", \"missing\"])\n",
+    "\n",
+    "        plot_idx += 1\n",
+    "\n",
+    "    # Hide unused subplots\n",
+    "    for i in range(plot_idx, len(axs)):\n",
+    "        fig.delaxes(axs[i])\n",
+    "\n",
+    "    # plt.tight_layout()\n",
+    "    plt.savefig(\"c2db-confusion_matrices.pdf\", bbox_inches=\"tight\")\n",
+    "    plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 52,
+   "id": "573b3c38",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "from mlip_arena.models import MLIPEnum\n",
+    "\n",
+    "thres = -1e-7\n",
+    "\n",
+    "summary_df = pd.DataFrame(columns=[\"Model\", \"Stable F1\", \"Unstable F1\", \"Weighted F1\"])\n",
+    "\n",
+    "for model in MLIPEnum:\n",
+    "    fpath = DATA_DIR / f\"{model.name}.parquet\"\n",
+    "\n",
+    "    if not fpath.exists() or model.name not in select_models:\n",
+    "        # print(f\"File {fpath} does not exist\")\n",
+    "        continue\n",
+    "    df = pd.read_parquet(fpath)\n",
+    "\n",
+    "    df[\"eigval_min\"] = df[\"eigenvalues\"].apply(\n",
+    "        lambda x: x.min() if np.isreal(x).all() else thres\n",
+    "    )\n",
+    "    df[\"freq_min\"] = df[\"frequencies\"].apply(\n",
+    "        lambda x: x.min() if np.isreal(x).all() else thres\n",
+    "    )\n",
+    "    df[\"dyn_stab\"] = ~np.logical_or(df[\"eigval_min\"] < thres, df[\"freq_min\"] < thres)\n",
+    "\n",
+    "    arg = np.argsort(uids)\n",
+    "    uids = np.array(uids)[arg]\n",
+    "    stabilities = stabilities[arg]\n",
+    "\n",
+    "    sorted_df = df[df[\"uid\"].isin(uids)].sort_values(by=\"uid\")\n",
+    "\n",
+    "    # sorted_df = sorted_df.reindex(uids).reset_index()\n",
+    "    sorted_df = sorted_df.set_index(\"uid\").reindex(uids)  # .loc[uids].reset_index()\n",
+    "\n",
+    "    sorted_df = sorted_df.loc[uids]\n",
+    "    # mask = ~np.logical_or(sorted_df['dyn_stab'].isna().values, stabilities == None)\n",
+    "    mask = ~(stabilities == None)\n",
+    "\n",
+    "    y_true = stabilities[mask].astype(\"int\")\n",
+    "    y_pred = sorted_df[\"dyn_stab\"][mask].fillna(-1).astype(\"int\")\n",
+    "    cm = confusion_matrix(y_true, y_pred, labels=[1, 0, -1])\n",
+    "    # print(model)\n",
+    "    # print(cm)\n",
+    "    # print(classification_report(y_true, y_pred, labels=[1, 0], target_names=['stable', 'unstable'], digits=3, output_dict=False))\n",
+    "\n",
+    "    report = classification_report(\n",
+    "        y_true,\n",
+    "        y_pred,\n",
+    "        labels=[1, 0],\n",
+    "        target_names=[\"stable\", \"unstable\"],\n",
+    "        digits=3,\n",
+    "        output_dict=True,\n",
+    "    )\n",
+    "\n",
+    "    summary_df = pd.concat(\n",
+    "        [\n",
+    "            summary_df,\n",
+    "            pd.DataFrame(\n",
+    "                [\n",
+    "                    {\n",
+    "                        \"Model\": model.name,\n",
+    "                        \"Stable F1\": report[\"stable\"][\"f1-score\"],\n",
+    "                        \"Unstable F1\": report[\"unstable\"][\"f1-score\"],\n",
+    "                        \"Macro F1\": report[\"macro avg\"][\"f1-score\"],\n",
+    "                        # 'Micro F1': report['micro avg']['f1-score'],\n",
+    "                        \"Weighted F1\": report[\"weighted avg\"][\"f1-score\"],\n",
+    "                    }\n",
+    "                ]\n",
+    "            ),\n",
+    "        ],\n",
+    "        ignore_index=True,\n",
+    "    )\n",
+    "\n",
+    "    # break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 85,
+   "id": "df660870",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "summary_df = summary_df.sort_values(by=[\"Macro F1\", \"Weighted F1\"], ascending=False)\n",
+    "summary_df.to_latex(\"c2db_summary_table.tex\", index=False, float_format=\"%.3f\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 103,
+   "id": "18f4a59b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from matplotlib import cm\n",
+    "\n",
+    "# Metrics and bar settings\n",
+    "metrics = [\"Stable F1\", \"Unstable F1\", \"Macro F1\", \"Weighted F1\"]\n",
+    "bar_width = 0.2\n",
+    "x = np.arange(len(summary_df))\n",
+    "\n",
+    "# Get Set2 colormap (as RGBA)\n",
+    "cmap = plt.get_cmap(\"tab20\")\n",
+    "colors = {metric: cmap(i) for i, metric in enumerate(metrics)}\n",
+    "\n",
+    "with plt.style.context(\"default\"):\n",
+    "    plt.rcParams.update(\n",
+    "        {\n",
+    "            \"font.size\": SMALL_SIZE,\n",
+    "            \"axes.titlesize\": MEDIUM_SIZE,\n",
+    "            \"axes.labelsize\": MEDIUM_SIZE,\n",
+    "            \"xtick.labelsize\": MEDIUM_SIZE,\n",
+    "            \"ytick.labelsize\": MEDIUM_SIZE,\n",
+    "            \"legend.fontsize\": SMALL_SIZE,\n",
+    "            \"figure.titlesize\": BIGGER_SIZE,\n",
+    "        }\n",
+    "    )\n",
+    "\n",
+    "    fig, ax = plt.subplots(figsize=(4, 3), layout=\"constrained\")\n",
+    "\n",
+    "    # Bar positions\n",
+    "    positions = {\n",
+    "        \"Stable F1\": x - 1.5 * bar_width,\n",
+    "        \"Unstable F1\": x - 0.5 * bar_width,\n",
+    "        \"Macro F1\": x + 0.5 * bar_width,\n",
+    "        \"Weighted F1\": x + 1.5 * bar_width,\n",
+    "    }\n",
+    "\n",
+    "    # Plot each metric with assigned color\n",
+    "    for metric, pos in positions.items():\n",
+    "        ax.bar(\n",
+    "            pos, summary_df[metric], width=bar_width, label=metric, color=colors[metric]\n",
+    "        )\n",
+    "\n",
+    "    ax.set_xlabel(\"Model\")\n",
+    "    ax.set_ylabel(\"F1 Score\")\n",
+    "    # ax.set_title('F1 Scores by Model and Class')\n",
+    "    ax.set_xticks(x)\n",
+    "    ax.set_xticklabels(summary_df[\"Model\"], rotation=45, ha=\"right\")\n",
+    "    ax.legend(ncols=2, bbox_to_anchor=(0.5, 1), loc=\"upper center\", fontsize=SMALL_SIZE)\n",
+    "    # ax.legend(ncols=2, fontsize=SMALL_SIZE)\n",
+    "    ax.spines[[\"top\", \"right\"]].set_visible(False)\n",
+    "    plt.tight_layout()\n",
+    "    plt.ylim(0, 0.9)\n",
+    "    plt.grid(axis=\"y\", linestyle=\"--\", alpha=0.6)\n",
+    "\n",
+    "    plt.savefig(\"c2db_f1_bar.pdf\", bbox_inches=\"tight\")\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c50f705",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mlip-arena",
+   "language": "python",
+   "name": "mlip-arena"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

benchmarks/c2db/c2db-confusion_matrices.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:463968f63e87ca0a7acd2e719cc481d0e3c5f5dd69ccf8f8659bddf6aa3b1e93
+size 21238

benchmarks/c2db/c2db-f1_bar.pdf

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

benchmarks/c2db/c2db.db

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

benchmarks/c2db/copy.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fdc16667361b10bfb032862d5d0610c242d75cb88f7f3883c43a406b245e991
+size 21349

benchmarks/c2db/run.py

@@ -0,0 +1,213 @@
+from itertools import product
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from dask.distributed import Client
+from dask_jobqueue import SLURMCluster
+from mlip_arena.models import MLIPEnum
+from mlip_arena.tasks import ELASTICITY, OPT, PHONON
+from mlip_arena.tasks.optimize import run as OPT
+from mlip_arena.tasks.utils import get_calculator
+from numpy import linalg as LA
+from prefect import flow, task
+from prefect_dask import DaskTaskRunner
+from tqdm.auto import tqdm
+
+from ase.db import connect
+
+select_models = [
+    "ALIGNN",
+    "CHGNet",
+    "M3GNet",
+    "MACE-MP(M)",
+    "MACE-MPA",
+    "MatterSim",
+    "ORBv2",
+    "SevenNet",
+]
+
+def elastic_tensor_to_voigt(C):
+    """
+    Convert a rank-4 (3x3x3x3) elastic tensor into a rank-2 (6x6) tensor using Voigt notation.
+
+    Parameters:
+    C (numpy.ndarray): A 3x3x3x3 elastic tensor.
+
+    Returns:
+    numpy.ndarray: A 6x6 elastic tensor in Voigt notation.
+    """
+    # voigt_map = {
+    #     (0, 0): 0, (1, 1): 1, (2, 2): 2,  # Normal components
+    #     (1, 2): 3, (2, 1): 3,  # Shear components
+    #     (0, 2): 4, (2, 0): 4,
+    #     (0, 1): 5, (1, 0): 5
+    # }
+    voigt_map = {
+        (0, 0): 0,
+        (1, 1): 1,
+        (2, 2): -1,  # Normal components
+        (1, 2): -1,
+        (2, 1): -1,  # Shear components
+        (0, 2): -1,
+        (2, 0): -1,
+        (0, 1): 2,
+        (1, 0): 2,
+    }
+
+    C_voigt = np.zeros((3, 3))
+
+    for i in range(3):
+        for j in range(3):
+            for k in range(3):
+                for l in range(3):
+                    alpha = voigt_map[(i, j)]
+                    beta = voigt_map[(k, l)]
+
+                    if alpha == -1 or beta == -1:
+                        continue
+
+                    factor = 1
+                    # if alpha in [3, 4, 5]:
+                    if alpha == 2:
+                        factor = factor * (2**0.5)
+                    if beta == 2:
+                        factor = factor * (2**0.5)
+
+                    C_voigt[alpha, beta] = C[i, j, k, l] * factor
+
+    return C_voigt
+
+
+# -
+
+
+@task
+def run_one(model, row):
+    if Path(f"{model.name}.pkl").exists():
+        df = pd.read_pickle(f"{model.name}.pkl")
+
+        # if row.key_value_pairs.get('uid', None) in df['uid'].unique():
+        #     pass
+    else:
+        df = pd.DataFrame(columns=["model", "uid", "eigenvalues", "frequencies"])
+
+    atoms = row.toatoms()
+    # print(data := row.key_value_pairs)
+
+    calc = get_calculator(model)
+
+    result_opt = OPT(
+        atoms,
+        calc,
+        optimizer="FIRE",
+        criterion=dict(fmax=0.05, steps=500),
+        symmetry=True,
+    )
+
+    atoms = result_opt["atoms"]
+
+    result_elastic = ELASTICITY(
+        atoms,
+        calc,
+        optimizer="FIRE",
+        criterion=dict(fmax=0.05, steps=500),
+        pre_relax=False,
+    )
+
+    elastic_tensor = elastic_tensor_to_voigt(result_elastic["elastic_tensor"])
+    eigenvalues, eigenvectors = LA.eig(elastic_tensor)
+
+    outdir = Path(f"{model.name}") / row.key_value_pairs.get(
+        "uid", atoms.get_chemical_formula()
+    )
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    np.savez(outdir / "elastic.npz", tensor=elastic_tensor, eigenvalues=eigenvalues)
+
+    result_phonon = PHONON(
+        atoms,
+        calc,
+        supercell_matrix=(2, 2, 1),
+        outdir=outdir,
+    )
+
+    frequencies = result_phonon["phonon"].get_frequencies(q=(0, 0, 0))
+
+    new_row = pd.DataFrame(
+        [
+            {
+                "model": model.name,
+                "uid": row.key_value_pairs.get("uid", None),
+                "eigenvalues": eigenvalues,
+                "frequencies": frequencies,
+            }
+        ]
+    )
+
+    df = pd.concat([df, new_row], ignore_index=True)
+    df.drop_duplicates(subset=["model", "uid"], keep="last", inplace=True)
+
+    df.to_pickle(f"{model.name}.pkl")
+
+
+@flow
+def run_all():
+    import random
+
+    random.seed(0)
+
+    futures = []
+    with connect("c2db.db") as db:
+        random_indices = random.sample(range(1, len(db) + 1), 1000)
+        for row, model in tqdm(
+            product(db.select(filter=lambda r: r["id"] in random_indices), MLIPEnum)
+        ):
+            if model.name not in select_models:
+                continue
+            future = run_one.submit(model, row)
+            futures.append(future)
+    return [f.result(raise_on_failure=False) for f in futures]
+
+
+# +
+
+
+if __name__ == "__main__":
+    nodes_per_alloc = 1
+    gpus_per_alloc = 1
+    ntasks = 1
+
+    cluster_kwargs = dict(
+        cores=1,
+        memory="64 GB",
+        processes=1,
+        shebang="#!/bin/bash",
+        account="matgen",
+        walltime="00:30:00",
+        # job_cpu=128,
+        job_mem="0",
+        job_script_prologue=[
+            "source ~/.bashrc",
+            "module load python",
+            "source activate /pscratch/sd/c/cyrusyc/.conda/dev",
+        ],
+        job_directives_skip=["-n", "--cpus-per-task", "-J"],
+        job_extra_directives=[
+            "-J c2db",
+            "-q regular",
+            f"-N {nodes_per_alloc}",
+            "-C gpu",
+            f"-G {gpus_per_alloc}",
+        ],
+    )
+
+    cluster = SLURMCluster(**cluster_kwargs)
+    print(cluster.job_script())
+    cluster.adapt(minimum_jobs=25, maximum_jobs=50)
+    client = Client(cluster)
+    # -
+
+    run_all.with_options(
+        task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True
+    )()

benchmarks/energy_conservation/run.py

@@ -0,0 +1,214 @@
+"""
+Task for running MD simulations and computing the differential entropy
+of the simulated structures with respect to a reference dataset.
+
+See https://github.com/dskoda/quests for differential entropy details.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+import numpy as np
+from ase.io import read
+from prefect import task
+from prefect.cache_policies import INPUTS, TASK_SOURCE
+from prefect.runtime import task_run
+
+from mlip_arena.models import MLIPEnum
+from mlip_arena.tasks.md import run as MD
+from mlip_arena.tasks.utils import logger
+
+try:
+    from quests.descriptor import get_descriptors
+    from quests.entropy import delta_entropy
+except ImportError as e:
+    logger.warning(e)
+    logger.warning(
+        "quests is not installed. Please install it using `pip install quests` or following the instructions at https://github.com/dskoda/quests to use this module."
+    )
+
+
+def get_entropy_from_path(
+    subset_path, dataset_path, dataset_desc_out_path, k=32, cutoff=5.0, h=0.015
+):
+    """
+    Computes the differential entropy of a subset of structures with respect
+    to a reference dataset.
+
+    Arguments:
+        subset_path (str): Path to the file containing the subset of structures.
+        dataset_path (str): Path to the file containing the full dataset of structures without the subset.
+        dataset_desc_out_path (str): Path to save the descriptors of the full dataset.
+        k (int, optional): Number of nearest neighbors used for descriptor calculation. Default is 32.
+        cutoff (float, optional): Cutoff distance for descriptor calculation. Default is 5.0.
+        h (float, optional): Bandwidth for the Gaussian kernel. Default is 0.015.
+
+    Returns:
+        np.ndarray: The differential entropy of the subset with respect to the dataset.
+    """
+
+    x_structures = read(dataset_path, index=":")
+    x_desc = get_descriptors(x_structures, k=k, cutoff=cutoff)
+    np.save(dataset_desc_out_path, x_desc)
+
+    y_structures = read(subset_path, index=":")
+    y_desc = get_descriptors(y_structures, k=k, cutoff=cutoff)
+
+    dH = delta_entropy(y_desc, x_desc, h=h)
+    return dH
+
+
+def get_trajectory_entropy(
+    trajectory_dir,
+    start_idx,
+    end_idx,
+    step,
+    dataset_desc_path,
+    k=32,
+    cutoff=5.0,
+    h=0.015,
+):
+    """
+    Computes the differential entropy of a subset of structures in a trajectory with respect
+    to a reference dataset.
+
+    Arguments:
+        trajectory_dir (str): Path to the directory containing the trajectory files.
+        start_idx (int): Starting index of the subset of structures to select from each trajectory.
+        end_idx (int): Ending index of the subset of structures to select from each trajectory.
+        step (int): Step size of the subset of structures to select from each trajectory.
+        dataset_desc_path (str): Path to the file containing the descriptors of the full dataset of structures without the subset.
+        k (int, optional): Number of nearest neighbors used for descriptor calculation. Default is 32.
+        cutoff (float, optional): Cutoff distance for descriptor calculation. Default is 5.0.
+        h (float, optional): Bandwidth for the Gaussian kernel. Default is 0.015.
+
+    Choose start_idx, end_idx, step to select which structures to compute the differential entropy for, based on what sliding window is chosen.
+    e.g. window of size 5 with stride 2 means we select every other structure starting at index 2 (middle of the first window) to the -2 index (middle of the last window)
+
+    Returns:
+        np.ndarray: The differential entropy of the subset of structures in the trajectory with respect to the dataset.
+    """
+    structures = []
+    for traj_file in sorted(os.listdir(trajectory_dir)):
+        traj = read(os.path.join(trajectory_dir, traj_file), index=":")
+        every_other = traj[start_idx:end_idx:step]
+        structures.extend(every_other)
+
+    desc = get_descriptors(structures, k=k, cutoff=cutoff)
+    x_desc = np.load(dataset_desc_path)
+    dH = delta_entropy(desc, x_desc, h=h)
+    return dH
+
+
+def run_simulations(model_names, structures, out_dir):
+    """
+    Runs simulations on a list of structures.
+
+    Parameters:
+        model_names (list[str]): List of models to use.
+        structures (list[ase.Atoms]): List of structures to simulate.
+        out_dir (str): Directory to save the simulation trajectories to.
+
+    Notes:
+        Structures are replicated to have at least 100 atoms and at most 500 atoms.
+        Structures are simulated with NVE MD at 1000 K for 5 ps.
+        Simulation trajectories are saved to files in out_dir, with each file named according to the index of the structure in the list.
+    """
+    min_atoms = 100
+    max_atoms = 500
+
+    futures = []
+
+    for model_name in model_names:
+        os.makedirs(out_dir, exist_ok=True)
+        model = MLIPEnum[model_name]
+        calc = model.value()
+
+        for i, atoms in enumerate(structures):
+            logger.info(
+                f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Running {model_name} on structure number {i}"
+            )
+
+            # Replicate the structure
+            n_atoms = len(atoms)
+            rep_factor = int(
+                np.ceil((min_atoms / n_atoms) ** (1 / 3))
+            )  # cube root since it's a 3D replication
+            supercell_atoms = atoms.repeat((rep_factor, rep_factor, rep_factor))
+            if len(supercell_atoms) > max_atoms:
+                logger.info(
+                    f"Skipping structure {i} because it has too many atoms ({len(supercell_atoms)} > {max_atoms})"
+                )
+                continue  # skip if it becomes too large
+
+            # Run NVE MD @ 1000K for 5 ps
+            future = MD.submit(
+                supercell_atoms,
+                calculator=calc,
+                ensemble="nve",
+                dynamics="velocityverlet",
+                time_step=1.0,  # fs
+                total_time=5000,  # 5 ps = 5000 fs
+                temperature=1000.0,
+                traj_file=f"{out_dir}/{i}.traj",
+                traj_interval=100,
+                zero_linear_momentum=True,
+                zero_angular_momentum=True,
+            )
+            futures.append(future)
+
+    return [f.result(raise_on_failure=False) for f in futures]
+
+
+def _generate_task_run_name():
+    task_name = task_run.task_name
+    parameters = task_run.parameters
+
+    trajectory_dir = parameters["trajectory_dir"]
+    dataset_desc_path = parameters["dataset_desc_path"]
+
+    return f"{task_name}: {trajectory_dir} - {dataset_desc_path}"
+
+
+@task(
+    name="Entropy along trajectory",
+    task_run_name=_generate_task_run_name,
+    cache_policy=TASK_SOURCE + INPUTS,
+)
+def run(
+    dataset_path,
+    model_names,
+    structures,
+    trajectory_dir,
+    start_idx,
+    end_idx,
+    step,
+    dataset_desc_path,
+    dH_out_path,
+    k=32,
+    cutoff=5.0,
+    h=0.015,
+):
+    # Get descriptors for the dataset. This should exclude the subset of structures used for simulations.
+    # This may take a while if the dataset is large - in that case, would recommend splitting the structures into separate chunks.
+    x_structures = read(dataset_path, index=":")
+    x_desc = get_descriptors(x_structures, k=k, cutoff=cutoff)
+    np.save(dataset_desc_path, x_desc)
+
+    # Run simulations
+    run_simulations(model_names, structures, trajectory_dir)
+
+    # Get entropy for structures along trajectories
+    dH = get_trajectory_entropy(
+        trajectory_dir,
+        start_idx,
+        end_idx,
+        step,
+        dataset_desc_path,
+        k=k,
+        cutoff=cutoff,
+        h=h,
+    )
+    np.save(dH_out_path, dH)

benchmarks/eos_bulk/CHGNet.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68871d694e93a3c3e7e272b9cbd87d3757e3bc689f30f3189db232d76e629c07
+size 429910

benchmarks/eos_bulk/CHGNet_processed.parquet

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

benchmarks/eos_bulk/M3GNet.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53dde465b5e10edd677f131f8a531e3dfc36303dd7ec7b9df0060c19847494d9
+size 427419

benchmarks/eos_bulk/M3GNet_processed.parquet

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

benchmarks/eos_bulk/MACE-MP(M).parquet

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

benchmarks/eos_bulk/MACE-MP(M)_processed.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e5507cdc5fe558b5d3fe2ea8f1dd577ac444e82c5347b5fbe738a4f855dffcb
+size 397379

benchmarks/eos_bulk/MACE-MPA.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53fcd188baddd4d5e797c5aa3de1b4368db711ebd29b7877cfe224856ba9d171
+size 428888

benchmarks/eos_bulk/MACE-MPA_processed.parquet

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

benchmarks/eos_bulk/MatterSim.parquet

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

benchmarks/eos_bulk/MatterSim_processed.parquet

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

benchmarks/eos_bulk/ORBv2.parquet

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

benchmarks/eos_bulk/ORBv2_processed.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7eb0a3060b8a2d3541b8fb1083176c88aae0a8be0008e84d5770998b01742216
+size 402554

benchmarks/eos_bulk/SevenNet.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64be88ec2632cdabf79daa01acb2cf2ef19fef0557813df5502c4f71ec566f4e
+size 428341

benchmarks/eos_bulk/SevenNet_processed.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9f484928f5086e8d1411a198ac69bfe44313597909b72c8676e9131cce1660f1
+size 398295

benchmarks/eos_bulk/analyze.py

@@ -0,0 +1,223 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from ase.db import connect
+from scipy import stats
+
+from mlip_arena.models import REGISTRY, MLIPEnum
+
+DATA_DIR = Path(__file__).parent.absolute()
+
+
+def load_wbm_structures():
+    """
+    Load the WBM structures from a ASE DB file.
+    """
+    with connect(DATA_DIR.parent / "wbm_structures.db") as db:
+        for row in db.select():
+            yield row.toatoms(add_additional_information=True)
+
+def gather_results():
+    for model in MLIPEnum:
+        if "eos_bulk" not in REGISTRY[model.name].get("gpu-tasks", []):
+            continue
+
+        if (DATA_DIR / f"{model.name}.parquet").exists():
+            continue
+
+        all_data = []
+
+        for atoms in load_wbm_structures():
+            fpath = Path(model.name) / f"{atoms.info['key_value_pairs']['wbm_id']}.pkl"
+            if not fpath.exists():
+                continue
+
+            all_data.append(pd.read_pickle(fpath))
+
+        df = pd.concat(all_data, ignore_index=True)
+        df.to_parquet(DATA_DIR / f"{model.name}.parquet")
+
+
+def summarize():
+    summary_table = pd.DataFrame(
+        columns=[
+            "model",
+            "energy-diff-flip-times",
+            "tortuosity",
+            "spearman-compression-energy",
+            "spearman-compression-derivative",
+            "spearman-tension-energy",
+            "missing",
+        ]
+    )
+
+
+    for model in MLIPEnum:
+        fpath = DATA_DIR / f"{model.name}.parquet"
+        if not fpath.exists():
+            continue
+        df_raw_results = pd.read_parquet(fpath)
+
+        df_analyzed = pd.DataFrame(
+            columns=[
+                "model",
+                "structure",
+                "formula",
+                "volume-ratio",
+                "energy-delta-per-atom",
+                "energy-diff-flip-times",
+                "energy-delta-per-volume-b0",
+                "tortuosity",
+                "spearman-compression-energy",
+                "spearman-compression-derivative",
+                "spearman-tension-energy",
+                "missing",
+            ]
+        )
+
+        for wbm_struct in load_wbm_structures():
+            structure_id = wbm_struct.info["key_value_pairs"]["wbm_id"]
+
+            try:
+                results = df_raw_results.loc[df_raw_results["id"] == structure_id]
+                b0 = results["b0"].values[0]
+                # vol0 = results["v0"].values[0]
+                results = results["eos"].values[0]
+                es = np.array(results["energies"])
+                vols = np.array(results["volumes"])
+
+                indices = np.argsort(vols)
+                vols = vols[indices]
+                es = es[indices]
+
+                imine = len(es) // 2
+                # min_center_val = np.min(es[imid - 1 : imid + 2])
+                # imine = np.where(es == min_center_val)[0][0]
+                emin = es[imine]
+                vol0 = vols[imine]
+
+                interpolated_volumes = [
+                    (vols[i] + vols[i + 1]) / 2 for i in range(len(vols) - 1)
+                ]
+                ediff = np.diff(es)
+                ediff_sign = np.sign(ediff)
+                mask = ediff_sign != 0
+                ediff = ediff[mask]
+                ediff_sign = ediff_sign[mask]
+                ediff_flip = np.diff(ediff_sign) != 0
+
+                etv = np.sum(np.abs(np.diff(es)))
+
+                data = {
+                    "model": model.name,
+                    "structure": structure_id,
+                    "formula": wbm_struct.get_chemical_formula(),
+                    "missing": False,
+                    "volume-ratio": vols / vol0,
+                    "energy-delta-per-atom": (es - emin) / len(wbm_struct),
+                    "energy-diff-flip-times": np.sum(ediff_flip).astype(int),
+                    "energy-delta-per-volume-b0": (es - emin) / (b0*vol0),
+                    "tortuosity": etv / (abs(es[0] - emin) + abs(es[-1] - emin)),
+                    "spearman-compression-energy": stats.spearmanr(
+                        vols[:imine], es[:imine]
+                    ).statistic,
+                    "spearman-compression-derivative": stats.spearmanr(
+                        interpolated_volumes[:imine], ediff[:imine]
+                    ).statistic,
+                    "spearman-tension-energy": stats.spearmanr(
+                        vols[imine:], es[imine:]
+                    ).statistic,
+                }
+
+            except Exception as e:
+                print(e)
+                data = {
+                    "model": model.name,
+                    "structure": structure_id,
+                    "formula": wbm_struct.get_chemical_formula(),
+                    "missing": True,
+                    "volume-ratio": None,
+                    "energy-delta-per-atom": None,
+                    "energy-delta-per-volume-b0": None,
+                    "energy-diff-flip-times": None,
+                    "tortuosity": None,
+                    "spearman-compression-energy": None,
+                    "spearman-compression-derivative": None,
+                    "spearman-tension-energy": None,
+                }
+
+            df_analyzed = pd.concat([df_analyzed, pd.DataFrame([data])], ignore_index=True)
+
+        df_analyzed.to_parquet(DATA_DIR / f"{model.name}_processed.parquet")
+        # json_fpath = DATA_DIR / f"EV_scan_analyzed_{model.name}.json"
+
+        # df_analyzed.to_json(json_fpath, orient="records")
+
+        valid_results = df_analyzed[df_analyzed["missing"] == False]
+
+        analysis_summary = {
+            "model": model.name,
+            "energy-diff-flip-times": valid_results["energy-diff-flip-times"].mean(),
+            "energy-diff-flip-times-std": valid_results["energy-diff-flip-times"].std(),
+            "tortuosity": valid_results["tortuosity"].mean(),
+            "tortuosity-std": valid_results["tortuosity"].std(),
+            "spearman-compression-energy": valid_results[
+                "spearman-compression-energy"
+            ].mean(),
+            "spearman-compression-energy-std": valid_results["spearman-compression-energy"].std(),
+            "spearman-compression-derivative": valid_results[
+                "spearman-compression-derivative"
+            ].mean(),
+            "spearman-compression-derivative-std": valid_results[
+                "spearman-compression-derivative"
+            ].std(),
+            "spearman-tension-energy": valid_results["spearman-tension-energy"].mean(),
+            "spearman-tension-energy-std": valid_results["spearman-tension-energy"].std(),
+            "missing": len(df_analyzed[df_analyzed["missing"] == True]),
+        }
+        summary_table = pd.concat(
+            [summary_table, pd.DataFrame([analysis_summary])], ignore_index=True
+        )
+
+
+    flip_rank = (
+        (summary_table["energy-diff-flip-times"] - 1)
+        .abs()
+        .rank(ascending=True, method="min")
+    )
+    tortuosity_rank = summary_table["tortuosity"].rank(ascending=True, method="min")
+    spearman_compression_energy_rank = summary_table["spearman-compression-energy"].rank(
+        method="min"
+    )
+    spearman_compression_derivative_rank = summary_table[
+        "spearman-compression-derivative"
+    ].rank(ascending=False, method="min")
+    spearman_tension_energy_rank = summary_table["spearman-tension-energy"].rank(
+        ascending=False, method="min"
+    )
+    missing_rank = summary_table["missing"].rank(ascending=True, method="min")
+
+    rank_aggr = (
+        flip_rank
+        + tortuosity_rank
+        + spearman_compression_energy_rank
+        + spearman_compression_derivative_rank
+        + spearman_tension_energy_rank
+        + missing_rank
+    )
+    rank = rank_aggr.rank(method="min")
+
+    summary_table.insert(1, "rank", rank.astype(int))
+    summary_table.insert(2, "rank-aggregation", rank_aggr.astype(int))
+    summary_table = summary_table.sort_values(by="rank", ascending=True)
+    summary_table = summary_table.reset_index(drop=True)
+
+    summary_table.to_csv(DATA_DIR / "summary.csv", index=False)
+    summary_table.to_latex(DATA_DIR / "summary.tex", index=False, float_format="%.3f")
+
+    return summary_table
+
+if __name__ == "__main__":
+    gather_results()
+    summarize()

benchmarks/eos_bulk/eSEN.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4503e17151b7376bbd88dc8c4767747e7290e8eae898e050b0a231a5c447e3e6
+size 427652

benchmarks/eos_bulk/eSEN_processed.parquet

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

benchmarks/eos_bulk/plot.py

@@ -0,0 +1,119 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from ase.db import connect
+
+from mlip_arena.models import REGISTRY as MODELS
+
+DATA_DIR = Path(__file__).parent.absolute()
+
+# Use a qualitative color palette from matplotlib
+palette_name = "tab10"  # Better for distinguishing multiple lines
+color_sequence = plt.get_cmap(palette_name).colors
+
+valid_models = [
+    model
+    for model, metadata in MODELS.items()
+    if "eos_bulk" in metadata.get("gpu-tasks", [])
+]
+
+def load_wbm_structures():
+    """
+    Load the WBM structures from a ASE DB file.
+    """
+    with connect(DATA_DIR.parent / "wbm_structures.db") as db:
+        for row in db.select():
+            yield row.toatoms(add_additional_information=True)
+
+# # Collect valid models first
+# valid_models = []
+# for model_name in valid_models:
+#     fpath = DATA_DIR / f"{model_name}_processed.parquet"
+#     if fpath.exists():
+#         df = pd.read_parquet(fpath)
+#         if len(df) > 0:
+#             valid_models.append(model)
+
+# # Ensure we're showing all 8 models
+# if len(valid_models) < 8:
+#     print(f"Warning: Only found {len(valid_models)} valid models instead of 8")
+
+# Set up the grid layout
+n_models = len(valid_models)
+n_cols = 4  # Use 4 columns
+n_rows = (n_models + n_cols - 1) // n_cols  # Ceiling division to get required rows
+
+# Create figure with enough space for all subplots
+fig = plt.figure(
+    figsize=(6, 1.25 * n_rows),  # Wider for better readability
+    constrained_layout=True,  # Better than tight_layout for this case
+)
+
+# Create grid of subplots
+axes = []
+for i in range(n_models):
+    ax = plt.subplot(n_rows, n_cols, i+1)
+    axes.append(ax)
+
+SMALL_SIZE = 6
+MEDIUM_SIZE = 8
+LARGE_SIZE = 10
+
+# Fill in the subplots with data
+for i, model_name in enumerate(valid_models):
+    fpath = DATA_DIR / f"{model_name}_processed.parquet"
+    df = pd.read_parquet(fpath)
+
+    ax = axes[i]
+    valid_structures = []
+
+    for j, (_, row) in enumerate(df.iterrows()):
+        structure_id = row["structure"]
+        formula = row.get("formula", "")
+        if isinstance(row["volume-ratio"], (list, np.ndarray)) and isinstance(
+            row["energy-delta-per-volume-b0"], (list, np.ndarray)
+        ):
+            vol_strain = row["volume-ratio"]
+            energy_delta = row["energy-delta-per-volume-b0"]
+            color = color_sequence[j % len(color_sequence)]
+            ax.plot(
+                vol_strain,
+                energy_delta,
+                color=color,
+                linewidth=1,
+                alpha=0.9,
+            )
+            valid_structures.append(structure_id)
+
+    # Set subplot title
+    ax.set_title(f"{model_name} ({len(valid_structures)})", fontsize=MEDIUM_SIZE)
+    
+    # Only add y-label to leftmost plots (those with index divisible by n_cols)
+    if i % n_cols == 0:
+        ax.set_ylabel("$\\frac{\\Delta E}{B V_0}$", fontsize=MEDIUM_SIZE)
+    else:
+        ax.set_ylabel("")
+    
+    # Only add x-label to bottom row plots
+    # Check if this plot is in the bottom row
+    is_bottom_row = (i // n_cols) == (n_rows - 1) or (i >= n_models - n_cols)
+    if is_bottom_row:
+        ax.set_xlabel("$V/V_0$", fontsize=MEDIUM_SIZE)
+    else:
+        ax.set_xlabel("")
+    
+    ax.set_ylim(-0.02, 0.1)  # Consistent y-limits
+    ax.axvline(x=1, linestyle="--", color="gray", alpha=0.7)
+    ax.tick_params(axis="both", which="major", labelsize=MEDIUM_SIZE)
+
+# Make sure all subplots share the x and y limits
+for ax in axes:
+    ax.set_xlim(0.8, 1.2)  # Adjust these as needed
+    ax.set_ylim(-0.02, 0.1)
+
+# Save the figure with all plots
+plt.savefig(DATA_DIR / "eos-bulk-grid.png", dpi=300, bbox_inches="tight")
+plt.savefig(DATA_DIR / "eos-bulk-grid.pdf", bbox_inches="tight")
+# plt.show()

benchmarks/eos_bulk/preprocessing.py

@@ -0,0 +1,12 @@
+import json
+
+from ase.db import connect
+from pymatgen.core import Structure
+
+with open("wbm_structures.json") as f:
+    structs = json.load(f)
+
+with connect("wbm_structures.db") as db:
+    for id, s in structs.items():
+        atoms = Structure.from_dict(s).to_ase_atoms(msonable=False)
+        db.write(atoms, wbm_id=id)

benchmarks/eos_bulk/run.py

@@ -0,0 +1,170 @@
+# import functools
+from pathlib import Path
+
+import pandas as pd
+from ase import Atoms
+from ase.db import connect
+from dask.distributed import Client
+from dask_jobqueue import SLURMCluster
+from prefect import flow, task
+from prefect.cache_policies import INPUTS, TASK_SOURCE
+from prefect.runtime import task_run
+from prefect_dask import DaskTaskRunner
+
+from mlip_arena.models import REGISTRY, MLIPEnum
+from mlip_arena.tasks.utils import get_calculator
+
+
+@task
+def load_wbm_structures():
+    """
+    Load the WBM structures from an ASE database file.
+    
+    Reads structures from 'wbm_structures.db' and yields them as ASE Atoms objects
+    with additional metadata preserved from the database.
+    
+    Yields:
+        ase.Atoms: Individual atomic structures from the WBM database with preserved
+                  metadata in the .info dictionary.
+    """
+    with connect("../wbm_structures.db") as db:
+        for row in db.select():
+            yield row.toatoms(add_additional_information=True)
+
+
+# def save_result(
+#     tsk: Task,
+#     run: TaskRun,
+#     state: State,
+#     model_name: str,
+#     id: str,
+# ):
+#     result = run.state.result()
+
+#     assert isinstance(result, dict)
+
+#     result["method"] = model_name
+#     result["id"] = id
+#     result.pop("atoms", None)
+
+#     fpath = Path(f"{model_name}")
+#     fpath.mkdir(exist_ok=True)
+
+#     fpath = fpath / f"{result['id']}.pkl"
+
+#     df = pd.DataFrame([result])
+#     df.to_pickle(fpath)
+
+
+@task(
+    name="EOS bulk - WBM",
+    task_run_name=lambda: f"{task_run.task_name}: {task_run.parameters['atoms'].get_chemical_formula()} - {task_run.parameters['model'].name}",
+    cache_policy=TASK_SOURCE + INPUTS,
+)
+def eos_bulk(atoms: Atoms, model: MLIPEnum):
+
+    from mlip_arena.tasks.eos import run as EOS
+    from mlip_arena.tasks.optimize import run as OPT
+
+    calculator = get_calculator(
+        model
+    )  # avoid sending entire model over prefect and select freer GPU
+
+    result = OPT.with_options(
+        refresh_cache=True,
+    )(
+        atoms,
+        calculator,
+        optimizer="FIRE",
+        criterion=dict(
+            fmax=0.1,
+        ),
+    )
+    result =  EOS.with_options(
+        refresh_cache=True,
+        # on_completion=[functools.partial(
+        #     save_result,
+        #     model_name=model.name,
+        #     id=atoms.info["key_value_pairs"]["wbm_id"],
+        # )],
+    )(
+        atoms=result["atoms"],
+        calculator=calculator,
+        optimizer="FIRE",
+        npoints=21,
+        max_abs_strain=0.2,
+        concurrent=False
+    )
+
+    result["method"] = model.name
+    result["id"] = atoms.info["key_value_pairs"]["wbm_id"]
+    result.pop("atoms", None)
+
+    fpath = Path(f"{model.name}")
+    fpath.mkdir(exist_ok=True)
+
+    fpath = fpath / f"{result['id']}.pkl"
+
+    df = pd.DataFrame([result])
+    df.to_pickle(fpath)
+
+    return df
+
+
+@flow
+def submit_tasks():
+    futures = []
+    for atoms in load_wbm_structures():
+        model = MLIPEnum["eSEN"]
+        # for model in MLIPEnum:
+        if "eos_bulk" not in REGISTRY[model.name].get("gpu-tasks", []):
+            continue
+        try:
+            result = eos_bulk.with_options(
+                refresh_cache=True
+            ).submit(atoms, model)
+            futures.append(result)
+        except Exception:
+            # print(f"Failed to submit task for {model.name}: {e}")
+            continue
+    return [f.result(raise_on_failure=False) for f in futures]
+
+
+if __name__ == "__main__":
+    nodes_per_alloc = 1
+    gpus_per_alloc = 1
+    ntasks = 1
+
+    cluster_kwargs = dict(
+        cores=1,
+        memory="64 GB",
+        shebang="#!/bin/bash",
+        account="m3828",
+        walltime="00:30:00",
+        job_mem="0",
+        job_script_prologue=[
+            "source ~/.bashrc",
+            "module load python",
+            "module load cudatoolkit/12.4",
+            "source activate /pscratch/sd/c/cyrusyc/.conda/dev",
+        ],
+        job_directives_skip=["-n", "--cpus-per-task", "-J"],
+        job_extra_directives=[
+            "-J eos_bulk",
+            "-q regular",
+            f"-N {nodes_per_alloc}",
+            "-C gpu",
+            f"-G {gpus_per_alloc}",
+            # "--exclusive",
+        ],
+    )
+
+    cluster = SLURMCluster(**cluster_kwargs)
+    print(cluster.job_script())
+    cluster.adapt(minimum_jobs=50, maximum_jobs=50)
+    client = Client(cluster)
+
+    submit_tasks.with_options(
+        task_runner=DaskTaskRunner(address=client.scheduler.address),
+        log_prints=True,
+    )()