update external calculator classes; update diatomic curves

.gitignore

@@ -1,10 +1,9 @@
-tests/
 *.out
 *.ipynb
 *.extxyz
 *.traj
 mlip_arena/tasks/*/*/*/
-
+lab/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

environment.yml

@@ -1,3 +1,4 @@
+name: mlip-arena
 channels:
   - defaults
   - conda-forge
@@ -14,7 +15,6 @@ dependencies:
   - libuuid=1.41.5=h5eee18b_0
   - ncurses=6.4=h6a678d5_0
   - openssl=3.0.13=h7f8727e_0
-  - pip=23.3.1=py311h06a4308_0
   - python=3.11.8=h955ad1f_0
   - readline=8.2=h5eee18b_0
   - setuptools=68.2.2=py311h06a4308_0
@@ -30,6 +30,7 @@ dependencies:
       - aioitertools==0.11.0
       - aiosignal==1.3.1
       - aiosqlite==0.20.0
+      - alabaster==0.7.16
       - alembic==1.13.1
       - alignn==2024.5.27
       - altair==5.3.0
@@ -37,26 +38,38 @@ dependencies:
       - anyio==3.7.1
       - appdirs==1.4.4
       - apprise==1.7.5
+      - argon2-cffi==23.1.0
+      - argon2-cffi-bindings==21.2.0
+      - arrow==1.3.0
       - ase==3.23.0
       - asgi-lifespan==2.1.0
       - asttokens==2.4.1
+      - async-lru==2.0.4
       - async-timeout==4.0.3
       - asyncpg==0.29.0
       - atomate2==0.0.14.post30+g256b39a1
       - attrs==23.2.0
       - autograd==1.5
+      - autopep8==2.3.1
       - autoray==0.6.9
+      - babel==2.15.0
       - bcrypt==4.1.2
+      - beautifulsoup4==4.12.3
       - bidict==0.23.1
+      - bleach==6.1.0
       - blinker==1.7.0
       - blosc2==2.7.0
+      - bokeh==2.4.3
+      - bokeh-sampledata==2024.2
       - boto3==1.34.74
       - botocore==1.34.74
+      - braceexpand==0.1.7
+      - cached-path==1.6.3
       - cachetools==5.3.3
-      - certifi==2024.7.4
+      - certifi==2024.8.30
       - cffi==1.16.0
       - charset-normalizer==3.3.2
-      - chgnet==0.3.5
+      - chgnet==0.4.0
       - click==8.1.7
       - cloudpickle==3.0.0
       - colorama==0.4.6
@@ -71,14 +84,21 @@ dependencies:
       - cython==3.0.10
       - dask==2024.3.1
       - dask-jobqueue==0.8.5
+      - datasets==2.21.0
       - dateparser==1.2.0
       - debugpy==1.8.1
       - decorator==5.1.1
-      - dgl==2.3.0+cu121
+      - defusedxml==0.7.1
+      - deprecation==2.1.0
+      - dgl==2.4.0+cu121
+      - dglgo==0.0.2
+      - dill==0.3.8
       - distributed==2024.3.1
+      - dm-tree==0.1.8
       - dnspython==2.6.1
       - docker==6.1.3
       - docker-pycreds==0.4.0
+      - docutils==0.21.2
       - e3nn==0.5.1
       - email-validator==2.1.1
       - emmet-core==0.82.1
@@ -86,19 +106,26 @@ dependencies:
       - fairchem-core==1.0.0
       - fastapi==0.110.0
       - fastjsonschema==2.20.0
-      - filelock==3.15.4
+      - filelock==3.16.1
       - fireworks==2.0.3
       - flake8==7.1.0
       - flask==3.0.2
       - flask-paginate==2024.3.28
       - fonttools==4.50.0
+      - fqdn==1.5.1
       - frozenlist==1.4.1
-      - fsspec==2024.6.1
+      - fsspec==2024.9.0
       - furl==2.1.3
       - future==1.0.0
       - gitdb==4.0.11
       - gitpython==3.1.43
+      - google-api-core==2.19.2
       - google-auth==2.29.0
+      - google-cloud-core==2.4.1
+      - google-cloud-storage==2.18.2
+      - google-crc32c==1.6.0
+      - google-resumable-media==2.7.2
+      - googleapis-common-protos==1.65.0
       - gpaw==24.7.0b1
       - greenlet==3.0.3
       - griffe==0.42.1
@@ -113,13 +140,18 @@ dependencies:
       - httpx==0.27.0
       - huggingface-hub==0.22.2
       - hyperframe==6.0.1
-      - idna==3.7
+      - icalendar==5.0.13
+      - idna==3.10
+      - imagesize==1.4.1
       - importlib-metadata==7.1.0
       - importlib-resources==6.1.3
       - inflect==7.3.1
       - ipykernel==6.29.4
+      - ipyspeck==0.6.1
       - ipython==8.22.2
-      - ipywidgets==8.1.2
+      - ipywidgets==8.1.3
+      - isoduration==20.11.0
+      - isort==5.13.2
       - itsdangerous==2.1.2
       - jarvis-tools==2024.5.10
       - jedi==0.19.1
@@ -127,17 +159,29 @@ dependencies:
       - jmespath==1.0.1
       - jobflow==0.1.17
       - joblib==1.3.2
+      - json5==0.9.25
       - jsonpatch==1.33
       - jsonpointer==2.4
       - jsonschema==4.21.1
       - jsonschema-specifications==2023.12.1
       - jupyter-client==8.6.1
       - jupyter-core==5.7.2
-      - jupyterlab-widgets==3.0.10
+      - jupyter-events==0.10.0
+      - jupyter-lsp==2.2.5
+      - jupyter-packaging==0.12.3
+      - jupyter-server==2.14.2
+      - jupyter-server-terminals==0.5.3
+      - jupyterlab==4.2.3
+      - jupyterlab-pygments==0.3.0
+      - jupyterlab-server==2.27.2
+      - jupyterlab-widgets==3.0.11
+      - kaleido==0.2.1
       - kiwisolver==1.4.5
       - kubernetes==29.0.0
       - latexcodec==3.0.0
+      - lightning==2.3.3
       - lightning-utilities==0.11.2
+      - littleutils==0.2.4
       - llvmlite==0.42.0
       - lmdb==1.4.1
       - lmdbm==0.0.5
@@ -149,30 +193,37 @@ dependencies:
       - markdown==3.6
       - markdown-it-py==2.2.0
       - markupsafe==2.1.5
-      - matgl==1.0.0
+      - matgl==1.1.2
       - matplotlib==3.8.3
       - matplotlib-inline==0.1.6
       - matscipy==1.0.0
       - mccabe==0.7.0
       - mdurl==0.1.2
-      - mlip-arena==0.0.1
+      - mistune==3.0.2
+      - mlip-arena==0.0.1a0
       - mongogrant==0.3.3
       - mongomock==4.1.2
-      - monty==2024.2.26
+      - monty==2024.7.30
       - more-itertools==10.3.0
       - mp-api==0.41.2
       - mpire==2.10.1
       - mpmath==1.3.0
       - msgpack==1.0.8
       - multidict==6.0.5
+      - multiprocess==0.70.16
       - natsort==8.4.0
+      - nbclient==0.10.0
+      - nbconvert==7.16.4
       - nbformat==5.10.4
       - ndindex==1.8
       - nest-asyncio==1.6.0
       - networkx==3.3
+      - notebook==7.2.1
+      - notebook-shim==0.2.4
       - numba==0.59.1
       - numexpr==2.10.1
       - numpy==1.26.4
+      - numpydoc==1.7.0
       - nvidia-cublas-cu12==12.1.3.1
       - nvidia-cuda-cupti-cu12==12.1.105
       - nvidia-cuda-nvrtc-cu12==12.1.105
@@ -183,38 +234,48 @@ dependencies:
       - nvidia-cusolver-cu12==11.4.5.107
       - nvidia-cusparse-cu12==12.1.0.106
       - nvidia-ml-py3==7.352.0
-      - nvidia-nccl-cu12==2.19.3
-      - nvidia-nvjitlink-cu12==12.5.82
+      - nvidia-nccl-cu12==2.20.5
+      - nvidia-nvjitlink-cu12==12.6.68
       - nvidia-nvtx-cu12==12.1.105
       - oauthlib==3.2.2
+      - ogb==1.3.6
       - opt-einsum==3.3.0
       - opt-einsum-fx==0.1.4
+      - orb-models==0.3.1
       - orderedmultidict==1.0.1
       - orjson==3.10.0
+      - outdated==0.2.2
+      - overrides==7.7.0
       - packaging==24.0
       - palettable==3.3.3
       - pandas==2.2.2
+      - pandocfilters==1.5.1
       - paramiko==3.4.0
       - parso==0.8.3
       - partd==1.4.1
       - pathspec==0.12.1
+      - patsy==0.5.6
       - pendulum==2.1.2
       - pennylane==0.32.0
       - pennylane-lightning==0.33.1
       - pexpect==4.9.0
       - pillow==10.2.0
+      - pip==24.2
       - platformdirs==4.2.0
       - plotly==5.20.0
       - plumed==2.9.0
       - prefect==2.16.8
       - prefect-dask==0.2.6
       - prettytable==3.10.0
+      - prometheus-client==0.20.0
       - prompt-toolkit==3.0.43
+      - proto-plus==1.24.0
       - protobuf==4.25.3
       - psutil==6.0.0
       - ptyprocess==0.7.0
       - pure-eval==0.2.2
       - py-cpuinfo==9.0.0
+      - py3dmol==2.0.0.post2
       - pyarrow==16.1.0
       - pyasn1==0.6.0
       - pyasn1-modules==0.4.0
@@ -229,15 +290,17 @@ dependencies:
       - pydocstyle==6.3.0
       - pyflakes==3.2.0
       - pygments==2.17.2
-      - pymatgen==2024.4.13
+      - pymatgen==2024.9.17.1
       - pymongo==4.6.3
       - pynacl==1.5.0
+      - pynanoflann==0.0.9
       - pyparsing==2.4.7
       - python-dateutil==2.9.0.post0
       - python-dotenv==1.0.1
       - python-engineio==4.9.0
       - python-graphviz==0.20.3
       - python-hostlist==1.23.0
+      - python-json-logger==2.0.7
       - python-multipart==0.0.9
       - python-slugify==8.0.4
       - python-socketio==5.11.2
@@ -246,12 +309,14 @@ dependencies:
       - pytzdata==2020.1
       - pyyaml==6.0.1
       - pyzmq==25.1.2
+      - rdkit-pypi==2022.9.5
       - readchar==4.0.6
       - referencing==0.34.0
       - regex==2023.12.25
       - requests==2.32.3
       - requests-oauthlib==2.0.0
       - rfc3339-validator==0.1.4
+      - rfc3986-validator==0.1.1
       - rich==13.3.5
       - rpds-py==0.18.0
       - rsa==4.9
@@ -260,12 +325,14 @@ dependencies:
       - rustworkx==0.14.2
       - s3transfer==0.10.1
       - safetensors==0.4.2
-      - scikit-learn==1.4.1.post1
-      - scipy==1.14.0
+      - scikit-learn==1.5.2
+      - scipy==1.14.1
       - semantic-version==2.10.0
+      - send2trash==1.8.3
       - sentinels==1.0.0
       - sentry-sdk==2.7.1
       - setproctitle==1.3.3
+      - sevenn==0.9.3.post1
       - shellingham==1.5.4
       - simple-websocket==1.0.0
       - simplejson==3.19.2
@@ -275,61 +342,80 @@ dependencies:
       - sniffio==1.3.1
       - snowballstemmer==2.2.0
       - sortedcontainers==2.4.0
-      - spglib==2.3.1
+      - soupsieve==2.5
+      - spglib==2.5.0
+      - sphinx==7.4.7
+      - sphinxcontrib-applehelp==1.0.8
+      - sphinxcontrib-devhelp==1.0.6
+      - sphinxcontrib-htmlhelp==2.0.6
+      - sphinxcontrib-jsmath==1.0.1
+      - sphinxcontrib-qthelp==1.0.8
+      - sphinxcontrib-serializinghtml==1.1.10
       - sqlalchemy==1.4.52
       - sqlalchemy-utils==0.41.2
       - sshtunnel==0.4.0
       - stack-data==0.6.3
       - starlette==0.36.3
+      - statsmodels==0.14.2
+      - stmol==0.0.9
       - streamlit==1.36.0
       - submitit==1.5.1
-      - sympy==1.12.1
+      - sympy==1.13.3
       - tables==3.9.2
       - tabulate==0.9.0
       - tblib==3.0.0
       - tenacity==8.2.3
       - tensorboard==2.17.0
       - tensorboard-data-server==0.7.2
+      - terminado==0.18.1
       - text-unidecode==1.3
       - threadpoolctl==3.4.0
+      - tinycss2==1.3.0
       - toml==0.10.2
+      - tomlkit==0.13.0
       - toolz==0.12.1
-      - torch==2.2.1
+      - torch==2.3.1
       - torch-dftd==0.4.0
       - torch-ema==0.3
       - torch-geometric==2.5.2
       - torch-scatter==2.1.2+pt22cu121
       - torch-sparse==0.6.18+pt22cu121
-      - torchdata==0.7.1
+      - torchdata==0.8.0
       - torchmetrics==1.3.2
       - tornado==6.4
-      - tqdm==4.66.4
+      - tqdm==4.66.5
       - trainstation==1.0
       - traitlets==5.14.2
-      - triton==2.2.0
+      - triton==2.3.1
       - typeguard==4.3.0
       - typer==0.12.0
       - typer-cli==0.12.0
       - typer-slim==0.12.0
+      - types-python-dateutil==2.9.0.20240316
       - typing-extensions==4.12.2
       - tzdata==2024.1
       - tzlocal==5.2
       - ujson==5.9.0
       - uncertainties==3.1.7
-      - urllib3==2.2.2
+      - uri-template==1.3.0
+      - urllib3==2.2.3
       - uvicorn==0.18.3
       - uvloop==0.19.0
       - wandb==0.17.4
       - watchdog==4.0.0
       - watchfiles==0.21.0
       - wcwidth==0.2.13
+      - webcolors==24.6.0
+      - webencodings==0.5.1
       - websocket-client==1.7.0
       - websockets==12.0
       - werkzeug==3.0.1
-      - widgetsnbextension==4.0.10
+      - widgetsnbextension==4.0.11
       - wrapt==1.16.0
       - wsproto==1.2.0
       - xmltodict==0.13.0
+      - xxhash==3.5.0
+      - xyzservices==2024.6.0
       - yarl==1.9.4
       - zict==3.0.0
       - zipp==3.18.1

mlip_arena/models/externals.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 import os
-import urllib
+from pathlib import Path
 from typing import Literal
 
 import matgl
+import requests
 import torch
 from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff
 from ase import Atoms
@@ -52,26 +55,28 @@ def get_freer_device() -> torch.device:
 
 
 class MACE_MP_Medium(MACECalculator):
-    def __init__(self, device=None, default_dtype="float32", **kwargs):
-        checkpoint_url = "http://tinyurl.com/5yyxdm76"
-        cache_dir = os.path.expanduser("~/.cache/mace")
+    def __init__(
+        self,
+        checkpoint="http://tinyurl.com/5yyxdm76",
+        device: str | None = None,
+        default_dtype="float32",
+        **kwargs,
+    ):
+        cache_dir = Path.home() / ".cache" / "mace"
         checkpoint_url_name = "".join(
-            c for c in os.path.basename(checkpoint_url) if c.isalnum() or c in "_"
+            c for c in os.path.basename(checkpoint) if c.isalnum() or c in "_"
         )
         cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
         if not os.path.isfile(cached_model_path):
+            import urllib
+
             os.makedirs(cache_dir, exist_ok=True)
-            # download and save to disk
-            print(f"Downloading MACE model from {checkpoint_url!r}")
-            _, http_msg = urllib.request.urlretrieve(checkpoint_url, cached_model_path)
+            _, http_msg = urllib.request.urlretrieve(checkpoint, cached_model_path)
             if "Content-Type: text/html" in http_msg:
                 raise RuntimeError(
-                    f"Model download failed, please check the URL {checkpoint_url}"
+                    f"Model download failed, please check the URL {checkpoint}"
                 )
-            print(f"Cached MACE model to {cached_model_path}")
         model = cached_model_path
-        msg = f"Using Materials Project MACE for MACECalculator with {model}"
-        print(msg)
 
         device = device or str(get_freer_device())
 
@@ -80,27 +85,30 @@ class MACE_MP_Medium(MACECalculator):
         )
 
 
+# TODO: could share the same class with MACE_MP_Medium
 class MACE_OFF_Medium(MACECalculator):
-    def __init__(self, device=None, default_dtype="float32", **kwargs):
-        checkpoint_url = "https://github.com/ACEsuit/mace-off/raw/main/mace_off23/MACE-OFF23_medium.model?raw=true"
-        cache_dir = os.path.expanduser("~/.cache/mace")
+    def __init__(
+        self,
+        checkpoint="https://github.com/ACEsuit/mace-off/raw/main/mace_off23/MACE-OFF23_medium.model?raw=true",
+        device: str | None = None,
+        default_dtype="float32",
+        **kwargs,
+    ):
+        cache_dir = Path.home() / ".cache" / "mace"
         checkpoint_url_name = "".join(
-            c for c in os.path.basename(checkpoint_url) if c.isalnum() or c in "_"
+            c for c in os.path.basename(checkpoint) if c.isalnum() or c in "_"
         )
         cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
         if not os.path.isfile(cached_model_path):
+            import urllib
+
             os.makedirs(cache_dir, exist_ok=True)
-            # download and save to disk
-            print(f"Downloading MACE model from {checkpoint_url!r}")
-            _, http_msg = urllib.request.urlretrieve(checkpoint_url, cached_model_path)
+            _, http_msg = urllib.request.urlretrieve(checkpoint, cached_model_path)
             if "Content-Type: text/html" in http_msg:
                 raise RuntimeError(
-                    f"Model download failed, please check the URL {checkpoint_url}"
+                    f"Model download failed, please check the URL {checkpoint}"
                 )
-            print(f"Cached MACE model to {cached_model_path}")
         model = cached_model_path
-        msg = f"Using Materials Project MACE for MACECalculator with {model}"
-        print(msg)
 
         device = device or str(get_freer_device())
 
@@ -112,15 +120,15 @@ class MACE_OFF_Medium(MACECalculator):
 class CHGNet(CHGNetCalculator):
     def __init__(
         self,
-        model: CHGNetModel | None = None,
-        use_device: str | None = None,
+        checkpoint: CHGNetModel | None = None,  # TODO: specifiy version
+        device: str | None = None,
         stress_weight: float | None = 1 / 160.21766208,
         on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
         **kwargs,
     ) -> None:
-        use_device = use_device or str(get_freer_device())
+        use_device = device or str(get_freer_device())
         super().__init__(
-            model=model,
+            model=checkpoint,
             use_device=use_device,
             stress_weight=stress_weight,
             on_isolated_atoms=on_isolated_atoms,
@@ -142,28 +150,31 @@ class CHGNet(CHGNetCalculator):
 class M3GNet(PESCalculator):
     def __init__(
         self,
+        checkpoint="M3GNet-MP-2021.2.8-PES",
+        # TODO: cannot assign device
         state_attr: torch.Tensor | None = None,
         stress_weight: float = 1.0,
         **kwargs,
     ) -> None:
-        potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+        potential = matgl.load_model(checkpoint)
         super().__init__(potential, state_attr, stress_weight, **kwargs)
 
 
 class EquiformerV2(OCPCalculator):
     def __init__(
         self,
-        model_name="EquiformerV2-lE4-lF100-S2EFS-OC22",
+        checkpoint="EquiformerV2-lE4-lF100-S2EFS-OC22",  # TODO: import from registry
+        # TODO: cannot assign device
         local_cache="/tmp/ocp/",
         cpu=False,
         seed=0,
         **kwargs,
     ) -> None:
         super().__init__(
-            model_name=model_name,
+            model_name=checkpoint,
             local_cache=local_cache,
             cpu=cpu,
-            seed=0,
+            seed=seed,
             **kwargs,
         )
 
@@ -178,17 +189,18 @@ class EquiformerV2(OCPCalculator):
 class EquiformerV2OC20(OCPCalculator):
     def __init__(
         self,
-        model_name="EquiformerV2-31M-S2EF-OC20-All+MD",
+        checkpoint="EquiformerV2-31M-S2EF-OC20-All+MD",  # TODO: import from registry
+        # TODO: cannot assign device
         local_cache="/tmp/ocp/",
         cpu=False,
         seed=0,
         **kwargs,
     ) -> None:
         super().__init__(
-            model_name=model_name,
+            model_name=checkpoint,
             local_cache=local_cache,
             cpu=cpu,
-            seed=0,
+            seed=seed,
             **kwargs,
         )
 
@@ -196,17 +208,18 @@ class EquiformerV2OC20(OCPCalculator):
 class eSCN(OCPCalculator):
     def __init__(
         self,
-        model_name="eSCN-L6-M3-Lay20-S2EF-OC20-All+MD",
+        checkpoint="eSCN-L6-M3-Lay20-S2EF-OC20-All+MD",  # TODO: import from registry
+        # TODO: cannot assign device
         local_cache="/tmp/ocp/",
         cpu=False,
         seed=0,
         **kwargs,
     ) -> None:
         super().__init__(
-            model_name=model_name,
+            model_name=checkpoint,
             local_cache=local_cache,
             cpu=cpu,
-            seed=0,
+            seed=seed,
             **kwargs,
         )
 
@@ -219,20 +232,50 @@ class eSCN(OCPCalculator):
 
 
 class ALIGNN(AlignnAtomwiseCalculator):
-    def __init__(self, dir_path: str = "/tmp/alignn/", device=None, **kwargs) -> None:
+    def __init__(self, device=None, dir_path: str = "/tmp/alignn/", **kwargs) -> None:
+        # TODO: cannot control version
         _ = get_figshare_model_ff(dir_path=dir_path)
         device = device or get_freer_device()
         super().__init__(path=dir_path, device=device, **kwargs)
 
 
 class SevenNet(SevenNetCalculator):
-    def __init__(self, device=None, **kwargs):
+    def __init__(
+        self,
+        checkpoint="7net-0",  # TODO: import from registry
+        device=None,
+        **kwargs,
+    ):
         device = device or get_freer_device()
-        super().__init__("7net-0", device=device, **kwargs)
+        super().__init__(checkpoint, device=device, **kwargs)
 
 
 class ORB(ORBCalculator):
-    def __init__(self, device=None, **kwargs):
+    def __init__(
+        self,
+        checkpoint="orbff-v1-20240827.ckpt",
+        device=None,
+        **kwargs,
+    ):
         device = device or get_freer_device()
-        orbff = pretrained.orb_v1(device=device)
+
+        cache_dir = Path.home() / ".cache" / "orb"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        ckpt_path = cache_dir / "orbff-v1-20240827.ckpt"
+
+        url = f"https://storage.googleapis.com/orbitalmaterials-public-models/forcefields/{checkpoint}"
+
+        if not ckpt_path.exists():
+            print(f"Downloading ORB model from {url} to {ckpt_path}...")
+            try:
+                response = requests.get(url, stream=True, timeout=120)
+                response.raise_for_status()
+                with open(ckpt_path, "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                print("Download completed.")
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError("Failed to download ORB model.") from e
+
+        orbff = pretrained.orb_v1(weights_path=ckpt_path, device=device)
         super().__init__(orbff, device=device, **kwargs)

mlip_arena/models/registry.yaml

@@ -27,7 +27,7 @@ CHGNet:
   module: externals
   class: CHGNet
   family: chgnet
-  package: chgnet==0.3.5
+  package: chgnet==0.3.8
   checkpoint:
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
@@ -66,6 +66,54 @@ M3GNet:
   nvt: true
   npt: true
 
+ORB:
+  module: externals
+  class: ORB
+  family: orb
+  package: orb-models==0.3.1
+  checkpoint: orbff-v1-20240827.ckpt
+  username: cyrusyc
+  last-update: 2024-03-25T14:30:00
+  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
+  datasets:
+    - atomind/mptrj # TODO: fake HF dataset repo
+    - atomind/alexandria
+  cpu-tasks:
+    - alexandria
+    - qmof
+  gpu-tasks:
+    - homonuclear-diatomics
+  github: https://github.com/orbital-materials/orb-models
+  doi: 
+  date: 2024-09-03
+  prediction: EFS
+  nvt: true
+  npt: true
+
+SevenNet:
+  module: externals
+  class: SevenNet
+  family: sevennet
+  package: sevenn==0.9.4
+  checkpoint: 7net-0
+  username: cyrusyc
+  last-update: 2024-03-25T14:30:00
+  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
+  datasets:
+    - atomind/mptrj # TODO: fake HF dataset repo
+  cpu-tasks:
+    - alexandria
+    - qmof
+  gpu-tasks:
+    - homonuclear-diatomics
+    - combustion
+  github: https://github.com/MDIL-SNU/SevenNet
+  doi: https://doi.org/10.1021/acs.jctc.4c00190
+  date: 2024-07-11
+  prediction: EFS
+  nvt: true
+  npt: true
+
 EquiformerV2(OC22):
   module: externals
   class: EquiformerV2
@@ -169,52 +217,4 @@ ALIGNN:
   npt: true
   github: https://github.com/usnistgov/alignn
   doi: https://doi.org/10.1038/s41524-021-00650-1
-  date: 2021-11-15
-
-SevenNet:
-  module: externals
-  class: SevenNet
-  family: sevennet
-  package: sevenn==0.9.4
-  checkpoint: 7net-0
-  username: cyrusyc
-  last-update: 2024-03-25T14:30:00
-  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
-  datasets:
-    - atomind/mptrj # TODO: fake HF dataset repo
-  cpu-tasks:
-    - alexandria
-    - qmof
-  gpu-tasks:
-    - homonuclear-diatomics
-    - combustion
-  github: https://github.com/MDIL-SNU/SevenNet
-  doi: https://doi.org/10.1021/acs.jctc.4c00190
-  date: 2024-07-11
-  prediction: EFS
-  nvt: true
-  npt: true
-
-ORB:
-  module: externals
-  class: ORB
-  family: orb
-  package: orb-models==0.3.1
-  checkpoint: orb_v1
-  username: cyrusyc
-  last-update: 2024-03-25T14:30:00
-  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
-  datasets:
-    - atomind/mptrj # TODO: fake HF dataset repo
-    - atomind/alexandria
-  cpu-tasks:
-    - alexandria
-    - qmof
-  gpu-tasks:
-    - homonuclear-diatomics
-  github: https://github.com/orbital-materials/orb-models
-  doi: 
-  date: 2024-09-03
-  prediction: EFS
-  nvt: true
-  npt: true
+  date: 2021-11-15

mlip_arena/tasks/combustion/H256O128.extxyz

@@ -0,0 +1,386 @@
+384
+Lattice="30.0 0.0 0.0 0.0 30.0 0.0 0.0 0.0 30.0" Properties=species:S:1:pos:R:3 Built=T with=T Packmol=T pbc="T T T"
+H        6.71669500      27.40330700      13.68784000
+H        6.56413900      26.94187400      13.13356600
+H        4.01874500      24.54731500       2.42614400
+H        4.06326800      24.60306800       1.69243900
+H       19.81990800      11.97306900      11.66638000
+H       19.80020800      12.70644500      11.59436300
+H       28.98607000      12.86336400      12.19767400
+H       29.00971400      12.88360500      12.93418200
+H       10.37632500       7.60336700       9.09607600
+H       10.09412100       7.01607100       9.44083300
+H       10.21969400      27.00047900      10.37176700
+H        9.56226800      26.66727600      10.35833100
+H        3.51642300       4.64618300      13.26100400
+H        3.19087000       4.01330300      13.06893100
+H       17.07272200      26.43633800       3.97721600
+H       17.76156200      26.60077000       4.18185100
+H        3.58254200       8.40601100       1.85885100
+H        3.32028600       8.99119200       1.49526000
+H       20.93526200      14.99894700      19.94272400
+H       20.92651600      14.99223100      20.67980700
+H       26.78606500      14.69290300      18.64689300
+H       26.43088200      14.16219600      19.01514500
+H       23.31077000      25.02192200      25.42458800
+H       23.23137700      24.98837100      26.15669800
+H        4.23870900      22.06352800       2.24193700
+H        4.13605700      21.89474500       2.95214000
+H       17.19183000      13.50635300      20.97020900
+H       17.33285700      13.55720300      20.24844800
+H       11.40128700      21.01758100      13.20015300
+H       11.23939200      21.24064200      13.88385500
+H        8.08553400      17.35614800       4.83985100
+H        7.85698800      18.00868300       4.58416100
+H       18.23851700      17.87971000      18.64072400
+H       18.18586700      17.93400600      19.37400000
+H        7.67448900       9.27045600       9.58857400
+H        8.38197500       9.28217700       9.38183700
+H       22.81810900       7.01891100      11.62916000
+H       23.02343500       6.88860700      10.93326000
+H        7.00678500      28.41965000      16.15063200
+H        6.61194800      29.00429800      16.36441700
+H       10.22790300      12.62515700      19.02217900
+H       10.06283300      12.03600000      19.43334700
+H        5.67633000      19.26890500       3.09689400
+H        6.16859600      19.35120500       3.63940200
+H        7.83689900      28.57325000       6.00799800
+H        8.20904000      29.01838100       5.55326500
+H       24.24356500       1.95484800      23.83583200
+H       24.33048700       1.23652000      23.97677000
+H        8.15609200      12.56884200       2.83826200
+H        8.47473700      12.51613700       2.17561500
+H        1.98944400      22.63517100       9.38706600
+H        1.99816600      22.52032700       8.65895300
+H        3.21198500      13.34481400      26.27085600
+H        2.83315800      12.78596400      26.56681200
+H       23.13687400      14.90616700      17.31808300
+H       23.60063200      15.47147100      17.41174900
+H       24.37760200      23.67694500       6.16965600
+H       24.26969800      23.54491000       5.45248300
+H       12.12534300       0.92505100       5.56915800
+H       12.59580200       1.08149500       6.11469100
+H       14.77057800      23.04450600      14.78597800
+H       14.88153400      22.32226700      14.88331100
+H       10.12077400       6.02995200      19.27346700
+H       10.56701000       5.58085700      19.65109300
+H        8.88571400      26.92992500      13.97009200
+H        9.13262800      27.42418800      13.48208400
+H       13.61412900      25.32932600       5.89573200
+H       14.17465000      25.02676400       5.52467300
+H       17.25064000       9.39374700      16.24924400
+H       17.23228900       9.42911200      16.98533300
+H       12.66108800      13.49816800       7.35969900
+H       12.75204000      13.66251500       6.64686500
+H       19.85540500      12.40447200      18.67913500
+H       19.87736300      13.08567100      18.39824500
+H       21.16039700      14.54404800      12.22178400
+H       20.77918300      14.69154600      11.60832300
+H        1.89709500      13.01516500      11.87415000
+H        1.61578300      12.78775700      11.23183900
+H       22.48727800      20.55335700      26.81374300
+H       22.08316200      19.93770500      26.84656000
+H       18.09641200      11.03175000      20.02888200
+H       18.08987100      11.14550300      19.30057500
+H        9.03487100      28.27705600      17.67592300
+H        9.02850300      28.47186100      18.38685500
+H       21.61332100      18.16863300      14.82530000
+H       20.91922800      18.33149900      15.01271300
+H        5.06777300       6.91012100      12.14248900
+H        5.14167800       7.53010400      12.53437400
+H       15.24240300       1.00156500       8.79894400
+H       15.15041500       1.70061000       8.58379800
+H        7.38156900      16.75420300      20.84490000
+H        7.49809200      17.42858000      20.57094300
+H        5.80395600       5.23791900      23.23164500
+H        5.59378400       5.32074900      23.93334300
+H       25.86193300      21.36196900      25.55815300
+H       26.43947000      21.14184700      25.15639300
+H       15.99876300      19.78997000       1.87456100
+H       15.41808700      19.79298200       2.32867600
+H       13.76283200      27.95478100       8.97671100
+H       13.72689900      27.99668400       8.24161500
+H        3.90503500      26.10922400      10.77631300
+H        3.36520200      25.83210000      11.19487700
+H        3.70175700       2.02893500      10.84822800
+H        3.98998200       2.68818200      11.00864200
+H       19.61345400      10.74535100      22.22723700
+H       19.65167500      10.72936900      22.96323800
+H       21.15335200      11.20594200      26.99475000
+H       21.53029400      10.81346000      27.49202900
+H       20.09434500       6.25424400      17.87873800
+H       20.01076700       6.21479100      18.61008800
+H        3.69537200      12.39851000       8.93646700
+H        4.09765200      11.81750500       9.14627000
+H       22.83093000      16.73250700      24.19709400
+H       23.31256100      17.17900100      23.86229600
+H        3.36198000      12.89766000      23.38943200
+H        3.45885500      13.00681500      24.11200700
+H        2.04611800      15.71757800       8.59329600
+H        1.43234600      15.77875600       8.18961900
+H        6.12945900      15.14026700      20.02887800
+H        6.02278200      14.57365100      20.48820400
+H        5.08911400      27.18139700       2.77337500
+H        5.06662100      27.18501900       3.51018800
+H       14.70758200       4.01177700      13.31696800
+H       14.89664300       4.24570600      12.64395500
+H        3.65814300      29.03356400      14.88286700
+H        3.73378000      28.95725100      14.15357400
+H       16.99951200      11.42586000      13.86881500
+H       17.03417100      11.12538100      13.19656200
+H        2.06797500      18.08136200      13.54583400
+H        1.48136600      18.09474000      13.99206700
+H       14.60250500       7.57683000       5.97130200
+H       15.33793200       7.61271900       6.00698800
+H       26.83711700       3.23048000       3.03884500
+H       26.55008700       3.90471800       3.11903900
+H       22.99777600      19.34731400      20.31560600
+H       22.45982800      19.77772100      20.57785400
+H       16.41966600       2.80599400      11.89128000
+H       16.56901900       2.08412200      11.88831400
+H       11.73017700      14.43811800      22.41741700
+H       11.12423500      14.82954800      22.56917700
+H       17.38249100      23.87139500      12.76358200
+H       17.38350900      23.82842900      12.02767000
+H        1.00061100      21.01224200       1.31470000
+H        1.44321900      21.03360900       1.90381400
+H        8.09004900       6.72936800       0.99999800
+H        8.19345600       7.35764200       1.37147100
+H        3.89822700      15.37766400       9.73104000
+H        4.54288400      15.70539800       9.87394700
+H       22.22552500      28.69576500      19.77030500
+H       21.65346500      29.00000300      19.41873800
+H        4.65573400      15.45943800      27.22827700
+H        5.35397400      15.56981400      27.43730400
+H       12.98406500       8.83317400      11.91500600
+H       12.47937800       9.06790600      11.43167900
+H       21.24968700       1.72985200      23.51877400
+H       21.91940400       2.02865000      23.44385400
+H       11.59386400      22.25408800      27.15665400
+H       11.73855200      21.54498900      27.29686000
+H       21.77510800       4.83204700      19.45360300
+H       21.74760500       4.86959100      18.71790800
+H       28.68304100       9.21677300      24.30293500
+H       28.79634600       9.49180700      23.62844900
+H       18.79159300       1.71351500       2.69284300
+H       18.77465700       2.18174600       3.26195300
+H       19.17586000      13.02896000      14.86997500
+H       19.71644200      13.51755200      14.75833800
+H        6.02966400      25.52593900      10.73507600
+H        5.80325300      25.22399900      10.10184400
+H        7.38506300       1.43983700      15.49859600
+H        6.79506700       1.06865500      15.73847300
+H       27.58382000      24.42158300      19.38106600
+H       27.42251200      25.04349100      19.01964600
+H        9.93092500      22.73321600      10.93929900
+H       10.22059100      22.24094300      11.40531900
+H       15.50084100      11.06308800       6.45439200
+H       15.87813400      11.34578600       5.88769500
+H       16.10016600      28.54356000       2.40068100
+H       16.19875700      27.92618400       2.01011700
+H       17.12622300       9.66589700       8.78315600
+H       16.48701000       9.96416700       8.99729500
+H       25.00042100      21.05189700      28.69399300
+H       25.35829500      20.90402200      28.06671900
+H       11.72392500      14.35901600      14.78824400
+H       11.48529200      14.43413500      14.09482800
+H       27.80088200      28.99951400      23.03485400
+H       27.35755700      28.41175300      23.07244700
+H        2.81120100      21.97327000       6.33202600
+H        3.02278700      22.64853600       6.53857400
+H       11.16091600      21.16245100      20.56587100
+H       10.97600400      21.57060200      19.98052200
+H       25.63314200       8.24689800      11.42146200
+H       25.84695500       7.98978900      10.76450500
+H        5.21035300      25.88590800      28.03985400
+H        4.61973200      25.44924700      28.10235800
+H       14.83558600      10.14445700      28.76566400
+H       14.45468700      10.77548600      28.75415200
+H       24.19325000      21.19802200      21.61625300
+H       24.60150400      21.42553100      22.18632500
+H        3.38383900      29.00000600       6.51473400
+H        3.38610800      28.83091700       7.23224200
+H       17.13203800      26.72357300      21.49248700
+H       17.08711600      27.41871500      21.25129500
+H       21.43007400      22.30567900      19.46406700
+H       20.99498700      21.73213800      19.30543700
+H       19.31644700      11.15485400      16.07821400
+H       19.68199100      11.18838700      16.71748400
+H       10.93924200       5.99489900      26.27955500
+H       10.41495800       5.47711900      26.30060500
+H       14.95008000      15.52212300       7.01442400
+H       14.72471900      15.69369600       7.69500400
+H       18.99727600      22.68120300      25.90377400
+H       19.40957700      23.25149200      25.68424600
+H        8.88447900       1.79024300      17.17555900
+H        8.86324600       1.77770600      17.91231200
+H       12.15121200       1.54535800      11.42100700
+H       11.71339400       1.67546500      11.99962700
+H        2.22952700      10.78313900      28.21835700
+H        2.22169400      10.63858800      27.49554500
+H        4.71118900      24.88843500       6.95444600
+H        4.75449800      25.06231100       6.23939000
+H        1.28060100       1.44032300      19.29701200
+H        1.83528100       1.36120500      18.81796700
+H       12.45357900       9.77922200      25.48060700
+H       11.74916000       9.99016300      25.53267400
+H        8.23385400      11.13248500      18.56147300
+H        8.05531600      10.83169400      19.21036700
+H       24.48638300      19.81957100      23.69282200
+H       25.13594000      19.58627200      23.95178100
+H       26.70376200       9.82401100      27.55103500
+H       26.90704500       9.33315600      28.06206400
+H        8.00838700       4.22559800      18.28897400
+H        8.42532000       3.95184600      17.74616600
+H       14.94452000       3.51127400       4.84402200
+H       14.89449900       3.42688600       5.57463200
+H       19.83606200       2.22162900      13.82277900
+H       20.02512800       2.59599500      13.21654700
+H        6.67903100       3.41714500       3.27561700
+H        5.96308500       3.29632200       3.40304000
+H        3.40442700       8.43161300      18.68125100
+H        3.47980300       8.38947800      19.41334200
+H       11.59437400      13.32368500      26.26415800
+H       12.23470100      13.62606800      26.46899700
+H       14.10053700       2.78998700       1.85965400
+H       14.21152000       2.83848100       2.58680300
+H       20.73807400      13.46815700      24.67708300
+H       21.38650500      13.80030400      24.56470100
+H       20.60974700      10.79901200      13.69614100
+H       20.72497000      10.71348400      14.41920500
+H        3.17495000      22.77164300      26.22314900
+H        3.62633300      22.80052800      25.64105600
+H       28.64961800      27.52039100      17.72169200
+H       28.65554600      27.60539900      18.45391600
+H        7.29178200      28.18007600       2.36813400
+H        7.28037500      28.22151000       3.10404600
+H        6.81956500      10.50511700       4.98278800
+H        6.16534900      10.84336000       4.95106300
+O       28.65774000      26.05310300       1.73187800
+O       27.82042800      25.18551200       1.41787100
+O        7.89599500      15.38902200      14.65745800
+O        7.88882100      15.60518500      13.43041800
+O       11.37704800      22.17970800       3.53597600
+O       10.62568400      23.16642300       3.65535900
+O       13.21571800      10.10612800       5.39037400
+O       13.81380800      10.31472700       4.31744300
+O       12.93699800      12.19119900      15.32963100
+O       13.36779300      11.62125600      16.35040800
+O       27.62157300       9.98656400       1.83057900
+O       26.68111300      10.54405800       1.23296700
+O        6.77274700      17.50780000       1.73044000
+O        5.79918400      17.87923300       1.04734400
+O       26.06597600      11.96086100      28.39447800
+O       26.32478700      12.75375600      27.46887600
+O        1.16753100       1.35249300      15.49177600
+O        1.18059100       1.20596000      14.25453600
+O       15.08210300      17.91068800      23.23158100
+O       15.17178200      18.46551400      22.11958700
+O       18.60542300      22.96513600      18.41436700
+O       19.71179100      23.45867400      18.70553300
+O        1.55709700      26.18859700       1.16408800
+O        2.61399600      26.49796400       1.74689400
+O       13.99292900      14.15190600      21.10355700
+O       14.04751700      14.18255600      19.85917400
+O       10.38541200      18.23372400      12.93027500
+O        9.91188600      17.49941400      12.04203600
+O       24.58403200      25.06294600      21.97592600
+O       25.05025500      26.20923800      21.83081500
+O       12.79523200      10.91568200      21.34176200
+O       12.59216100      11.07454900      22.56074900
+O        7.88339000      21.12603300       4.34091800
+O        8.23639900      20.09896500       3.73024400
+O       14.29669600      11.48202900      10.12026500
+O       14.65658800      11.78505600       8.96655000
+O       24.45593800      10.09229400      18.89934500
+O       23.27212000      10.31487400      19.21783900
+O       10.14103800      28.74148400       8.55297800
+O        9.09630200      28.47275800       9.17645200
+O        6.53750100      18.89849400      13.13483400
+O        7.18188100      19.75253500      12.49624700
+O        5.99854200      21.81797900       5.25108200
+O        5.99224800      23.06318800       5.20840100
+O        9.03156900       6.67285800      22.94269500
+O        7.94989000       6.26883900      23.41082700
+O       11.37806100      15.86060900      17.89431700
+O       10.58973300      15.69880900      16.94312500
+O       18.69252800      21.49332500      14.52289300
+O       19.04391300      21.05544000      13.41060100
+O       10.88775200      25.62526500      15.90813900
+O       10.56171400      25.66761900      17.10993500
+O        5.83323700       6.64702500      10.08718000
+O        5.78880400       6.30641500       8.88950900
+O        7.71126200       8.59122800      18.86977200
+O        7.73147600       8.60812400      17.62409400
+O       20.74284900       2.12287500      21.37159100
+O       20.54629600       2.47368100      20.19230800
+O        9.33194100      10.52852700      15.25052100
+O        8.99841100       9.39518300      14.85467600
+O       15.28732400      14.68259500       9.86397700
+O       14.24814200      15.19621500      10.32081500
+O        8.79509800       2.36395800      13.47414000
+O        8.02117200       2.64244400      12.53825000
+O       22.07041200      27.82260300      11.86453400
+O       22.47004200      27.29076900      10.81103700
+O       25.64900400      10.71424900      24.59372700
+O       25.14662500      10.86773900      25.72353300
+O       17.46952600      13.14388400      17.59581300
+O       17.02295500      12.33034400      16.76446800
+O       12.27907200       3.72064200       4.88171600
+O       12.34240900       3.77544700       3.63857800
+O       19.50544800      19.26652500      26.41877500
+O       19.46811300      19.42938700      27.65347700
+O       16.09394700      23.74355500       8.66223700
+O       16.43146400      24.61823500       7.84161300
+O       12.88585100      27.95040100      27.15844000
+O       13.33429000      28.28825200      28.27072000
+O       17.90265700      22.35366300       6.91603700
+O       18.19323300      21.88298700       5.79959800
+O       23.38376600       4.75998600       8.13216000
+O       23.39868600       4.98386700       9.35774600
+O       16.75294800       1.89397000       5.04381400
+O       17.32240900       2.99348800       5.18231300
+O       17.56519600       7.26680100      21.77071000
+O       18.31705700       8.25843700      21.83209000
+O       11.99835900      11.23426800      18.12324800
+O       12.36029500      11.59229700      19.26044800
+O        7.58959900      13.32747400      11.00503600
+O        6.37218600      13.53282600      10.83727300
+O       18.59862600      24.35057600       9.30150100
+O       18.94728500      24.62991400       8.13839600
+O       23.48008200      18.54852600      17.99906900
+O       23.36944900      17.83172600      16.98597300
+O       22.77455400      28.18135800      14.92714700
+O       23.86580000      27.65620500      14.63422500
+O       20.78948500       3.67260900       1.83948100
+O       21.56215500       3.51027100       2.80334500
+O        9.05637000      15.36340100      27.31637100
+O       10.18282500      14.88236800      27.54465200
+O        5.36124000      14.96264300      23.52556700
+O        4.24760100      14.87097200      22.97437700
+O       22.64266600      13.26344800      26.40413500
+O       22.64331400      14.50641200      26.49043800
+O        1.32756400      12.17561700      20.06774300
+O        2.10349400      12.35107600      21.02667600
+O        1.18281700       3.80763100       4.94719100
+O        1.01891800       2.82667000       5.69769700
+O       12.59019600      17.94027500      24.45964400
+O       12.50958200      18.08831700      23.22514300
+O        9.77496800      21.51131100       7.49102600
+O       10.45743600      20.53412700       7.85400700
+O        1.69369700       6.90651400       3.72886700
+O        1.19082500       6.11102100       4.54539400
+O       26.88872500      10.54198400      11.12473400
+O       27.32702000      10.28414700       9.98727000
+O        1.72700700      26.73221500      14.58446600
+O        2.56724200      25.91897900      15.01465200
+O       10.15694600      18.69534200      19.15618300
+O        9.99999700      19.29683100      20.23599100
+O        1.99644200      10.71276900      12.34430700
+O        2.09478800      10.26022600      13.50100000
+O        2.90497500       4.19153100       2.82131400
+O        1.84632400       4.84452600       2.74877200
+O        5.05591600      16.45745000      12.34283200
+O        6.13192000      15.97608900      11.93920500
+O        6.49320100      16.95653500      24.96253500
+O        6.59211000      17.17298400      26.18555300

mlip_arena/tasks/combustion/water.ipynb

@@ -119,7 +119,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -130,18 +130,29 @@
       "\n",
       "#SBATCH -A matgen\n",
       "#SBATCH --mem=0\n",
-      "#SBATCH -t 02:00:00\n",
+      "#SBATCH -t 01:45:00\n",
       "#SBATCH -J combustion-water\n",
       "#SBATCH -q regular\n",
       "#SBATCH -N 1\n",
       "#SBATCH -C gpu\n",
       "#SBATCH -G 4\n",
+      "#SBATCH --exclusive\n",
       "source ~/.bashrc\n",
       "module load python\n",
       "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.16:40657 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.22:36743 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
       "\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.\n",
+      "Perhaps you already have a cluster running?\n",
+      "Hosting the HTTP server on port 34167 instead\n",
+      "  warnings.warn(\n"
+     ]
     }
    ],
    "source": [
@@ -154,7 +165,7 @@
     "    memory=\"64 GB\",\n",
     "    shebang=\"#!/bin/bash\",\n",
     "    account=\"matgen\",\n",
-    "    walltime=\"02:00:00\",\n",
+    "    walltime=\"01:45:00\",\n",
     "    job_mem=\"0\",\n",
     "    job_script_prologue=[\n",
     "        \"source ~/.bashrc\",\n",
@@ -168,7 +179,7 @@
     "        f\"-N {nodes_per_alloc}\",\n",
     "        \"-C gpu\",\n",
     "        f\"-G {gpus_per_alloc}\",\n",
-    "        # f\"--exclusive\",\n",
+    "        f\"--exclusive\",\n",
     "        # \"--time-min=00:30:00\",\n",
     "        # \"--comment=1-00:00:00\",\n",
     "        # \"--signal=B:USR1@60\",\n",
@@ -187,7 +198,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -209,7 +220,7 @@
     "            pressure=None,\n",
     "            mb_velocity_seed=0,\n",
     "            traj_file=Path(REGISTRY[model.name][\"family\"])\n",
-    "            / f\"{atoms.get_chemical_formula()}.traj\",\n",
+    "            / f\"{model.name}_{atoms.get_chemical_formula()}.traj\",\n",
     "            traj_interval=1000,\n",
     "            restart=True,\n",
     "        )\n",
@@ -229,11 +240,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:32.556 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.187 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:32.556 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'airborne-chupacabra'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
+       "18:24:16.187 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'cinnamon-swine'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -242,11 +253,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:32.563 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/fb8616ef-5579-49a7-81ce-0306180ca401</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.205 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/b40f12fb-6644-4319-8d19-5e01e7e282aa</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:32.563 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/fb8616ef-5579-49a7-81ce-0306180ca401\u001b[0m\n"
+       "18:24:16.205 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/b40f12fb-6644-4319-8d19-5e01e7e282aa\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -255,11 +266,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:32.565 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.16:40657\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.207 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.22:36743\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:32.565 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.16:40657\n"
+       "18:24:16.207 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.22:36743\n"
       ]
      },
      "metadata": {},
@@ -268,11 +279,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:32.591 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - The Dask dashboard is available at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">http://128.55.64.16:8787/status</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.213 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - The Dask dashboard is available at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">http://128.55.64.22:34167/status</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:32.591 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.16:8787/status\u001b[0m\n"
+       "18:24:16.213 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.22:34167/status\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -281,11 +292,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.239 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-4' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.764 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-0' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.239 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-4' for task 'md'\n"
+       "18:24:16.764 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-0' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -294,11 +305,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.280 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-7' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.793 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-1' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.280 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-7' for task 'md'\n"
+       "18:24:16.793 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-1' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -307,11 +318,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.286 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-1' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.797 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-3' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.286 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-1' for task 'md'\n"
+       "18:24:16.797 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-3' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -320,11 +331,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.289 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-2' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.800 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-4' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.289 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-2' for task 'md'\n"
+       "18:24:16.800 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-4' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -333,11 +344,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.292 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-5' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.803 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-8' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.292 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-5' for task 'md'\n"
+       "18:24:16.803 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-8' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -346,11 +357,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.295 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-6' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.805 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-9' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.295 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-6' for task 'md'\n"
+       "18:24:16.805 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-9' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -359,11 +370,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.298 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-9' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.808 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-2' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.298 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-9' for task 'md'\n"
+       "18:24:16.808 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-2' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -372,11 +383,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.301 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-0' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.810 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-6' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.301 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-0' for task 'md'\n"
+       "18:24:16.810 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-6' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -385,11 +396,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.304 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-3' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.813 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-7' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.304 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-3' for task 'md'\n"
+       "18:24:16.813 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-7' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -398,11 +409,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:33.306 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Created task run 'md-8' for task 'md'\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:16.949 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Created task run 'md-5' for task 'md'\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:33.306 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Created task run 'md-8' for task 'md'\n"
+       "18:24:16.949 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Created task run 'md-5' for task 'md'\n"
       ]
      },
      "metadata": {},
@@ -411,11 +422,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.559 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-8' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.496 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-9' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.559 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-8' for execution.\n"
+       "18:24:17.496 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-9' for execution.\n"
       ]
      },
      "metadata": {},
@@ -424,11 +435,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.712 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-5' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.510 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-7' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.712 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-5' for execution.\n"
+       "18:24:17.510 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-7' for execution.\n"
       ]
      },
      "metadata": {},
@@ -437,11 +448,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.731 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-4' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.520 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-1' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.731 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-4' for execution.\n"
+       "18:24:17.520 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-1' for execution.\n"
       ]
      },
      "metadata": {},
@@ -450,11 +461,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.745 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-7' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.529 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-5' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.745 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-7' for execution.\n"
+       "18:24:17.529 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-5' for execution.\n"
       ]
      },
      "metadata": {},
@@ -463,11 +474,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.755 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-1' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.536 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-8' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.755 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-1' for execution.\n"
+       "18:24:17.536 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-8' for execution.\n"
       ]
      },
      "metadata": {},
@@ -476,11 +487,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.771 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-3' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.550 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-0' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.771 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-3' for execution.\n"
+       "18:24:17.550 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-0' for execution.\n"
       ]
      },
      "metadata": {},
@@ -489,11 +500,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.789 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-0' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.605 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-3' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.789 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-0' for execution.\n"
+       "18:24:17.605 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-3' for execution.\n"
       ]
      },
      "metadata": {},
@@ -502,11 +513,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.803 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-2' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.614 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-4' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.803 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-2' for execution.\n"
+       "18:24:17.614 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-4' for execution.\n"
       ]
      },
      "metadata": {},
@@ -515,11 +526,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.813 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-6' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.621 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-6' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.813 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-6' for execution.\n"
+       "18:24:17.621 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-6' for execution.\n"
       ]
      },
      "metadata": {},
@@ -528,11 +539,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:37.843 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'airborne-chupacabra'</span> - Submitted task run 'md-9' for execution.\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">18:24:17.635 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'cinnamon-swine'</span> - Submitted task run 'md-2' for execution.\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "18:24:37.843 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-9' for execution.\n"
+       "18:24:17.635 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'cinnamon-swine'\u001b[0m - Submitted task run 'md-2' for execution.\n"
       ]
      },
      "metadata": {},

mlip_arena/tasks/diatomics/alignn/run.ipynb

@@ -1,243 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "3200850a-b8fb-4f50-9815-16ae8da0f942",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "from ase import Atoms, Atom\n",
-    "from ase.io import read, write\n",
-    "from ase.data import chemical_symbols, covalent_radii, vdw_alvarez\n",
-    "from ase.parallel import paropen as open\n",
-    "\n",
-    "from pathlib import Path\n",
-    "import os\n",
-    "import numpy as np\n",
-    "from pymatgen.core import Element\n",
-    "from tqdm.auto import tqdm\n",
-    "import pandas as pd\n",
-    "\n",
-    "\n",
-    "from alignn.ff.ff import AlignnAtomwiseCalculator,default_path\n",
-    "\n",
-    "\n",
-    "model_path = default_path()\n",
-    "calc = AlignnAtomwiseCalculator(path=model_path, device='cuda')\n",
-    "\n",
-    "model_name = 'ALIGNN'\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "90887faa-1601-4c4c-9c44-d16731471d7f",
-   "metadata": {
-    "scrolled": true,
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "\n",
-    "\n",
-    "for symbol in tqdm(chemical_symbols):\n",
-    "    \n",
-    "    s = set([symbol])\n",
-    "    \n",
-    "    if 'X' in s:\n",
-    "        continue\n",
-    "        \n",
-    "    try:\n",
-    "        atom = Atom(symbol)\n",
-    "        rmin = covalent_radii[atom.number] * 0.95\n",
-    "        rvdw = vdw_alvarez.vdw_radii[atom.number] if atom.number < len(vdw_alvarez.vdw_radii) else np.nan \n",
-    "        rmax = 3.1 * rvdw if not np.isnan(rvdw) else 6\n",
-    "        rstep = 0.01 #if rmin < 1 else 0.4\n",
-    "\n",
-    "        a = 2 * rmax\n",
-    "\n",
-    "        npts = int((rmax - rmin)/rstep)\n",
-    "\n",
-    "        rs = np.linspace(rmin, rmax, npts)\n",
-    "        e = np.zeros_like(rs)\n",
-    "\n",
-    "        da = symbol + symbol\n",
-    "\n",
-    "        out_dir = Path(str(da))\n",
-    "\n",
-    "        os.makedirs(out_dir, exist_ok=True)\n",
-    "\n",
-    "        skip = 0\n",
-    "        \n",
-    "        element = Element(symbol)\n",
-    "        \n",
-    "        try:\n",
-    "            m = element.valence[1]\n",
-    "            if element.valence == (0, 2):\n",
-    "                m = 0\n",
-    "        except:\n",
-    "            m = 0\n",
-    "            \n",
-    "            \n",
-    "        r = rs[0]\n",
-    "        \n",
-    "        positions = [\n",
-    "            [a/2-r/2, a/2, a/2],\n",
-    "            [a/2+r/2, a/2, a/2],\n",
-    "        ]\n",
-    "            \n",
-    "        traj_fpath = out_dir / f\"{model_name}.extxyz\"\n",
-    "\n",
-    "        if traj_fpath.exists():\n",
-    "            traj = read(traj_fpath, index=\":\")\n",
-    "            skip = len(traj)\n",
-    "            atoms = traj[-1]\n",
-    "        else:\n",
-    "            # Create the unit cell with two atoms\n",
-    "            atoms = Atoms(\n",
-    "                da, \n",
-    "                positions=positions,\n",
-    "                # magmoms=magmoms,\n",
-    "                cell=[a, a+0.001, a+0.002], \n",
-    "                pbc=True\n",
-    "            )\n",
-    "        \n",
-    "        print(atoms)\n",
-    "\n",
-    "        calc = calc\n",
-    "\n",
-    "        atoms.calc = calc\n",
-    "\n",
-    "        for i, r in enumerate(tqdm(rs)):\n",
-    "\n",
-    "            if i < skip:\n",
-    "                continue\n",
-    "\n",
-    "            positions = [\n",
-    "                [a/2-r/2, a/2, a/2],\n",
-    "                [a/2+r/2, a/2, a/2],\n",
-    "            ]\n",
-    "            \n",
-    "            # atoms.set_initial_magnetic_moments(magmoms)\n",
-    "                \n",
-    "            atoms.set_positions(positions)\n",
-    "\n",
-    "            e[i] = atoms.get_potential_energy()\n",
-    "            \n",
-    "            atoms.calc.results.update({\n",
-    "                \"forces\": atoms.get_forces()\n",
-    "            })\n",
-    "\n",
-    "            write(traj_fpath, atoms, append=\"a\")\n",
-    "    except Exception as e:\n",
-    "        print(e)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a0ac2c09-370b-4fdd-bf74-ea5c4ade0215",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "\n",
-    "df = pd.DataFrame(columns=['name', 'method', 'R', 'E', 'F', 'S^2'])\n",
-    "\n",
-    "for symbol in tqdm(chemical_symbols):\n",
-    "    \n",
-    "    da = symbol + symbol\n",
-    "    \n",
-    "    out_dir = Path(da)\n",
-    "    \n",
-    "    traj_fpath = out_dir / f\"{model_name}.extxyz\"\n",
-    "\n",
-    "\n",
-    "    if traj_fpath.exists():\n",
-    "        traj = read(traj_fpath, index=\":\")\n",
-    "    else:\n",
-    "        continue\n",
-    "    \n",
-    "    Rs, Es, Fs, S2s = [], [], [], []\n",
-    "    for atoms in traj:\n",
-    "        \n",
-    "        vec = atoms.positions[1] - atoms.positions[0]\n",
-    "        r = np.linalg.norm(vec)\n",
-    "        e = atoms.get_potential_energy()\n",
-    "        f = np.inner(vec/r, atoms.get_forces()[1])\n",
-    "        # s2 = np.mean(np.power(atoms.get_magnetic_moments(), 2))\n",
-    "        \n",
-    "        Rs.append(r)\n",
-    "        Es.append(e)\n",
-    "        Fs.append(f)\n",
-    "        # S2s.append(s2)\n",
-    "        \n",
-    "    data = {\n",
-    "        'name': da,\n",
-    "        'method': 'ALIGNN',\n",
-    "        'R': Rs,\n",
-    "        'E': Es,\n",
-    "        'F': Fs,\n",
-    "        'S^2': S2s\n",
-    "    }\n",
-    "\n",
-    "    df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
-    "\n",
-    "json_fpath = 'homonuclear-diatomics.json'\n",
-    "\n",
-    "df.to_json(json_fpath, orient='records')    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e0dd4367-3dca-440f-a7a9-7fdd84183f2c",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "df"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4e6ae884-89f3-43f2-8fd9-19bf00c91566",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "mlip-arena",
-   "language": "python",
-   "name": "mlip-arena"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.8"
-  },
-  "widgets": {
-   "application/vnd.jupyter.widget-state+json": {
-    "state": {},
-    "version_major": 2,
-    "version_minor": 0
-   }
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

pyproject.toml

@@ -37,16 +37,37 @@ dependencies=[
 ]
 
 [project.optional-dependencies]
-m3gnet = ["matgl", "dgl", "torch<=2.2.1"]
-mace = ["mace-torch"]
-chgnet = ["chgnet"]
-fairchem = ["fairchem"]
+run = [
+    "torch==2.2.0",
+    "e3nn==0.5.1",
+    "matgl==1.1.2",
+    "dgl==2.4.0+cu121",
+    "mace-torch==0.3.4",
+    "chgnet==0.3.8",
+    "fairchem-core==0.1.0",
+    "sevenn==0.9.3.post1",
+    "orb-models==0.3.1",
+    "alignn==2024.5.27"
+]
 app = [
     "streamlit",
     "plotly",
     "bokeh==2.4.3",
     "statsmodels"
 ]
+test = [
+    "torch==2.2.0",
+    "e3nn==0.5.1",
+    "matgl==1.1.2",
+    "dgl==2.4.0+cu121",
+    "mace-torch==0.3.4",
+    "chgnet==0.3.8",
+    "fairchem-core==0.1.0",
+    "sevenn==0.9.3.post1",
+    "orb-models==0.3.1",
+    "alignn==2024.5.27",
+    "pytest"
+]
 
 [project.urls]
 Homepage = "https://github.com/atomind-ai/mlip-arena"

scripts/install-pyg.sh

@@ -1,15 +1,11 @@
 
 
 # PyTorch Geometric (OCP)
-TORCH=2.2.0
+TORCH=2.3.1
 CUDA=cu121
 
 pip install --verbose --no-cache torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
 pip install --verbose --no-cache torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
 
-# DGL (M3GNet)
-pip install --verbose --no-cache dgl -f https://data.dgl.ai/wheels/{CUDA}/repo.html
-
-
-# DGL (ALIGNN)
-# pip install --verbose --no-cache dgl -f https://data.dgl.ai/wheels/torch-2.2/cu122/repo.html
+# DGL (M3GNet, ALIGNN)
+pip install dgl -U -f https://data.dgl.ai/wheels/torch-2.3/cu121/repo.html

tests/download_models.py

@@ -1,5 +0,0 @@
-from huggingface_hub import hf_hub_download
-
-fpath = hf_hub_download(repo_id="cyrusyc/mace-universal", subfolder="pretrained", filename="2023-12-12-mace-128-L1_epoch-199.model")
-
-print(fpath)

tests/hf_hub.ipynb

@@ -1,480 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "from huggingface_hub import hf_hub_download\n",
-    "from ase.calculators.calculator import Calculator\n",
-    "# from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP\n",
-    "\n",
-    "from mlip_arena.models.externals import MACE_MP_Medium\n",
-    "\n",
-    "from mlip_arena.models.utils import MLIPMap, MLIPEnum"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 9,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "issubclass(MLIPEnum[\"MACE-MP(M)\"].value, Calculator)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "isinstance(MLIPEnum[\"MACE-MP(M)\"], MLIPEnum)# in MLIPEnum"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['MACE-MP(M)', 'CHGNet', 'EquiformerV2(OC22)', 'eSCN(OC20)']"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "MLIPEnum._member_names_"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'MACE-MP(M)': mlip_arena.models.externals.MACE_MP_Medium,\n",
-       " 'CHGNet': mlip_arena.models.externals.CHGNet,\n",
-       " 'EquiformerV2(OC22)': mlip_arena.models.externals.EquiformerV2,\n",
-       " 'eSCN(OC20)': mlip_arena.models.externals.eSCN}"
-      ]
-     },
-     "execution_count": 9,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "MLIPMap"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "MLIPEnum.MACE-MP(M)\n",
-      "MLIPEnum.CHGNet\n",
-      "MLIPEnum.EquiformerV2(OC22)\n",
-      "MLIPEnum.eSCN(OC20)\n"
-     ]
-    }
-   ],
-   "source": [
-    "for mlip in MLIPEnum:\n",
-    "    print(mlip)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "mlip_arena.models.externals.MACE_MP_Medium"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "MLIPMap['MACE-MP(M)']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Using Materials Project MACE for MACECalculator with /global/homes/c/cyrusyc/.cache/mace/5yyxdm76\n",
-      "Selected GPU cuda:0 with 40338.06 MB free memory from 1 GPUs\n",
-      "Default dtype float32 does not match model dtype float64, converting models to float32.\n"
-     ]
-    }
-   ],
-   "source": [
-    "mace_mp = MACE_MP_Medium()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Select GPU cuda:0 with 40316.98 MB free memory from 1 GPUs\n",
-      "CHGNet v0.3.0 initialized with 412,525 parameters\n",
-      "CHGNet will run on cuda:0\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "WARNING:root:Detected old config, converting to new format. Consider updating to avoid potential incompatibilities.\n",
-      "WARNING:root:Skipping scheduler setup. No training set found.\n"
-     ]
-    }
-   ],
-   "source": [
-    "from mlip_arena.models.externals import EquiformerV2, CHGNet\n",
-    "\n",
-    "chgnet = CHGNet()\n",
-    "\n",
-    "equiformer_v2 = EquiformerV2()\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "fpath = hf_hub_download(\n",
-    "    repo_id=\"cyrusyc/mace-universal\",\n",
-    "    subfolder=\"pretrained\",\n",
-    "    filename=\"2023-12-12-mace-128-L1_epoch-199.model\",\n",
-    "    revision=None,  # TODO: Add revision\n",
-    ")\n",
-    "\n",
-    "model = torch.load(fpath, map_location=\"cpu\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "module = ModuleMLIP(model=model)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "CommitInfo(commit_url='https://huggingface.co/atomind/mace-mp-medium/commit/eb12c5387b9e655d83a4e2e10c0f0779c3745227', commit_message='Push model using huggingface_hub.', commit_description='', oid='eb12c5387b9e655d83a4e2e10c0f0779c3745227', pr_url=None, pr_revision=None, pr_num=None)"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "module.save_pretrained(\n",
-    "    \"mace\",\n",
-    "    repo_id=\"atomind/MACE_MP_Medium\".lower().replace(\"_\", \"-\"),\n",
-    "    push_to_hub=True\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    }
-   ],
-   "source": [
-    "\n",
-    "from mlip_arena.models.mace import MACE_MP_Medium\n",
-    "import torch\n",
-    "\n",
-    "calc = MACE_MP_Medium(device=torch.device(\"cuda\"))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "ScaleShiftMACE(\n",
-       "  (node_embedding): LinearNodeEmbeddingBlock(\n",
-       "    (linear): Linear(89x0e -> 128x0e | 11392 weights)\n",
-       "  )\n",
-       "  (radial_embedding): RadialEmbeddingBlock(\n",
-       "    (bessel_fn): BesselBasis(r_max=6.0, num_basis=10, trainable=False)\n",
-       "    (cutoff_fn): PolynomialCutoff(p=5.0, r_max=6.0)\n",
-       "  )\n",
-       "  (spherical_harmonics): SphericalHarmonics()\n",
-       "  (atomic_energies_fn): AtomicEnergiesBlock(energies=[-3.6672, -1.3321, -3.4821, -4.7367, -7.7249, -8.4056, -7.3601, -7.2846, -4.8965, 0.0000, -2.7594, -2.8140, -4.8469, -7.6948, -6.9633, -4.6726, -2.8117, -0.0626, -2.6176, -5.3905, -7.8858, -10.2684, -8.6651, -9.2331, -8.3050, -7.0490, -5.5774, -5.1727, -3.2521, -1.2902, -3.5271, -4.7085, -3.9765, -3.8862, -2.5185, 6.7669, -2.5635, -4.9380, -10.1498, -11.8469, -12.1389, -8.7917, -8.7869, -7.7809, -6.8500, -4.8910, -2.0634, -0.6396, -2.7887, -3.8186, -3.5871, -2.8804, -1.6356, 9.8467, -2.7653, -4.9910, -8.9337, -8.7356, -8.0190, -8.2515, -7.5917, -8.1697, -13.5927, -18.5175, -7.6474, -8.1230, -7.6078, -6.8503, -7.8269, -3.5848, -7.4554, -12.7963, -14.1081, -9.3549, -11.3875, -9.6219, -7.3244, -5.3047, -2.3801, 0.2495, -2.3240, -3.7300, -3.4388, -5.0629, -11.0246, -12.2656, -13.8556, -14.9331, -15.2828])\n",
-       "  (interactions): ModuleList(\n",
-       "    (0): RealAgnosticResidualInteractionBlock(\n",
-       "      (linear_up): Linear(128x0e -> 128x0e | 16384 weights)\n",
-       "      (conv_tp): TensorProduct(128x0e x 1x0e+1x1o+1x2e+1x3o -> 128x0e+128x1o+128x2e+128x3o | 512 paths | 512 weights)\n",
-       "      (conv_tp_weights): FullyConnectedNet[10, 64, 64, 64, 512]\n",
-       "      (linear): Linear(128x0e+128x1o+128x2e+128x3o -> 128x0e+128x1o+128x2e+128x3o | 65536 weights)\n",
-       "      (skip_tp): FullyConnectedTensorProduct(128x0e x 89x0e -> 128x0e+128x1o | 1458176 paths | 1458176 weights)\n",
-       "      (reshape): reshape_irreps()\n",
-       "    )\n",
-       "    (1): RealAgnosticResidualInteractionBlock(\n",
-       "      (linear_up): Linear(128x0e+128x1o -> 128x0e+128x1o | 32768 weights)\n",
-       "      (conv_tp): TensorProduct(128x0e+128x1o x 1x0e+1x1o+1x2e+1x3o -> 256x0e+384x1o+384x2e+256x3o | 1280 paths | 1280 weights)\n",
-       "      (conv_tp_weights): FullyConnectedNet[10, 64, 64, 64, 1280]\n",
-       "      (linear): Linear(256x0e+384x1o+384x2e+256x3o -> 128x0e+128x1o+128x2e+128x3o | 163840 weights)\n",
-       "      (skip_tp): FullyConnectedTensorProduct(128x0e+128x1o x 89x0e -> 128x0e | 1458176 paths | 1458176 weights)\n",
-       "      (reshape): reshape_irreps()\n",
-       "    )\n",
-       "  )\n",
-       "  (products): ModuleList(\n",
-       "    (0): EquivariantProductBasisBlock(\n",
-       "      (symmetric_contractions): SymmetricContraction(\n",
-       "        (contractions): ModuleList(\n",
-       "          (0): Contraction(\n",
-       "            (contractions_weighting): ModuleList(\n",
-       "              (0-1): 2 x GraphModule()\n",
-       "            )\n",
-       "            (contractions_features): ModuleList(\n",
-       "              (0-1): 2 x GraphModule()\n",
-       "            )\n",
-       "            (weights): ParameterList(\n",
-       "                (0): Parameter containing: [torch.float64 of size 89x4x128 (cuda:0)]\n",
-       "                (1): Parameter containing: [torch.float64 of size 89x1x128 (cuda:0)]\n",
-       "            )\n",
-       "            (graph_opt_main): GraphModule()\n",
-       "          )\n",
-       "          (1): Contraction(\n",
-       "            (contractions_weighting): ModuleList(\n",
-       "              (0-1): 2 x GraphModule()\n",
-       "            )\n",
-       "            (contractions_features): ModuleList(\n",
-       "              (0-1): 2 x GraphModule()\n",
-       "            )\n",
-       "            (weights): ParameterList(\n",
-       "                (0): Parameter containing: [torch.float64 of size 89x6x128 (cuda:0)]\n",
-       "                (1): Parameter containing: [torch.float64 of size 89x1x128 (cuda:0)]\n",
-       "            )\n",
-       "            (graph_opt_main): GraphModule()\n",
-       "          )\n",
-       "        )\n",
-       "      )\n",
-       "      (linear): Linear(128x0e+128x1o -> 128x0e+128x1o | 32768 weights)\n",
-       "    )\n",
-       "    (1): EquivariantProductBasisBlock(\n",
-       "      (symmetric_contractions): SymmetricContraction(\n",
-       "        (contractions): ModuleList(\n",
-       "          (0): Contraction(\n",
-       "            (contractions_weighting): ModuleList(\n",
-       "              (0-1): 2 x GraphModule()\n",
-       "            )\n",
-       "            (contractions_features): ModuleList(\n",
-       "              (0-1): 2 x GraphModule()\n",
-       "            )\n",
-       "            (weights): ParameterList(\n",
-       "                (0): Parameter containing: [torch.float64 of size 89x4x128 (cuda:0)]\n",
-       "                (1): Parameter containing: [torch.float64 of size 89x1x128 (cuda:0)]\n",
-       "            )\n",
-       "            (graph_opt_main): GraphModule()\n",
-       "          )\n",
-       "        )\n",
-       "      )\n",
-       "      (linear): Linear(128x0e -> 128x0e | 16384 weights)\n",
-       "    )\n",
-       "  )\n",
-       "  (readouts): ModuleList(\n",
-       "    (0): LinearReadoutBlock(\n",
-       "      (linear): Linear(128x0e+128x1o -> 1x0e | 128 weights)\n",
-       "    )\n",
-       "    (1): NonLinearReadoutBlock(\n",
-       "      (linear_1): Linear(128x0e -> 16x0e | 2048 weights)\n",
-       "      (non_linearity): Activation [x] (16x0e -> 16x0e)\n",
-       "      (linear_2): Linear(16x0e -> 1x0e | 16 weights)\n",
-       "    )\n",
-       "  )\n",
-       "  (scale_shift): ScaleShiftBlock(scale=0.804154, shift=0.164097)\n",
-       ")"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "calc.model\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from mlip_arena.models import MLIP\n",
-    "\n",
-    "model = MLIP.from_pretrained(\"atomind/mace-mp-medium\", map_location=\"cuda\", revision=\"main\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<generator object Module.modules at 0x7ff33915f920>"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model.modules()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "AttributeError",
-     "evalue": "MLIP has no attribute `model`",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_submodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmodel\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/nn/modules/module.py:681\u001b[0m, in \u001b[0;36mModule.get_submodule\u001b[0;34m(self, target)\u001b[0m\n\u001b[1;32m    678\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m atoms:\n\u001b[1;32m    680\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(mod, item):\n\u001b[0;32m--> 681\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(mod\u001b[38;5;241m.\u001b[39m_get_name() \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m has no \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    682\u001b[0m                              \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mattribute `\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m item \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    684\u001b[0m     mod \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(mod, item)\n\u001b[1;32m    686\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(mod, torch\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mModule):\n",
-      "\u001b[0;31mAttributeError\u001b[0m: MLIP has no attribute `model`"
-     ]
-    }
-   ],
-   "source": [
-    "model.get_submodule(\"model\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "for name, param in model.named_parameters():\n",
-    "    print(name, param.data)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(module)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.8"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}