add orb model to diatomic tasks

CITATION.cff

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

mlip_arena/models/externals.py

@@ -3,15 +3,16 @@ import urllib
 from typing import Literal
 
 import matgl
-import requests
 import torch
-from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff, default_path
+from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff
 from ase import Atoms
 from chgnet.model.dynamics import CHGNetCalculator
-from chgnet.model.model import CHGNet
+from chgnet.model.model import CHGNet as CHGNetModel
 from fairchem.core import OCPCalculator
 from mace.calculators import MACECalculator
 from matgl.ext.ase import PESCalculator
+from orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
 from sevenn.sevennet_calculator import SevenNetCalculator
 
 
@@ -111,7 +112,7 @@ class MACE_OFF_Medium(MACECalculator):
 class CHGNet(CHGNetCalculator):
     def __init__(
         self,
-        model: CHGNet | None = None,
+        model: CHGNetModel | None = None,
         use_device: str | None = None,
         stress_weight: float | None = 1 / 160.21766208,
         on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
@@ -173,6 +174,7 @@ class EquiformerV2(OCPCalculator):
             force=atoms.get_forces(),
         )
 
+
 class EquiformerV2OC20(OCPCalculator):
     def __init__(
         self,
@@ -190,6 +192,7 @@ class EquiformerV2OC20(OCPCalculator):
             **kwargs,
         )
 
+
 class eSCN(OCPCalculator):
     def __init__(
         self,
@@ -217,25 +220,19 @@ class eSCN(OCPCalculator):
 
 class ALIGNN(AlignnAtomwiseCalculator):
     def __init__(self, dir_path: str = "/tmp/alignn/", device=None, **kwargs) -> None:
-        model_path = get_figshare_model_ff(dir_path=dir_path)
+        _ = get_figshare_model_ff(dir_path=dir_path)
         device = device or get_freer_device()
         super().__init__(path=dir_path, device=device, **kwargs)
 
-    def calculate(self, atoms, properties=None, system_changes=None):
-        super().calculate(atoms, properties, system_changes)
-
 
 class SevenNet(SevenNetCalculator):
     def __init__(self, device=None, **kwargs):
-        # url = (
-        #     "https://github.com/MDIL-SNU/SevenNet/raw/main/pretrained_potentials"
-        #     "/SevenNet_0__11July2024/checkpoint_sevennet_0.pth"
-        # )
-        # ckpt_cache = "/tmp/sevennet_checkpoint.pth.tar"
-        # response = requests.get(url)
-        # with open(ckpt_cache, mode="wb") as file:
-        #     file.write(response.content)
-
         device = device or get_freer_device()
-
         super().__init__("7net-0", device=device, **kwargs)
+
+
+class ORB(ORBCalculator):
+    def __init__(self, device=None, **kwargs):
+        device = device or get_freer_device()
+        orbff = pretrained.orb_v1(device=device)
+        super().__init__(orbff, device=device, **kwargs)

mlip_arena/models/registry.yaml

@@ -177,4 +177,26 @@ SevenNet:
   date: 2024-07-11
   prediction: EFS
   nvt: true
+  npt: true
+
+ORB:
+  module: externals
+  class: ORB
+  family: orb
+  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

mlip_arena/tasks/combustion/water.ipynb

@@ -139,7 +139,7 @@
       "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.32:35601 --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.16:40657 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
       "\n"
      ]
     }
@@ -229,11 +229,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\">10:53:24.417 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.417 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'masterful-panther'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -242,11 +242,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\">10:53:24.422 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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/4d0a29f8-4313-4dd3-b311-ea64abf82a83</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.422 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\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/4d0a29f8-4313-4dd3-b311-ea64abf82a83\u001b[0m\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -255,11 +255,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\">10:53:24.423 | <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.32:35601\n",
+       "<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>\n"
       ],
       "text/plain": [
-       "10:53:24.423 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.32:35601\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -268,11 +268,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\">10:53:24.430 | <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.32: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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.430 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.32:8787/status\u001b[0m\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -281,11 +281,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\">10:53:24.927 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.927 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-4' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -294,11 +294,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\">10:53:24.975 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.975 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-6' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -307,11 +307,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\">10:53:24.981 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.981 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-0' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -320,11 +320,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\">10:53:24.989 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.989 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-1' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -333,11 +333,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\">10:53:24.995 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.995 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-5' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -346,11 +346,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\">10:53:24.999 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:24.999 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-7' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -359,11 +359,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\">10:53:25.002 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.002 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-8' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -372,11 +372,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\">10:53:25.005 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.005 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-3' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -385,11 +385,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\">10:53:25.008 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.008 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-2' for task 'md'\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -398,11 +398,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\">10:53:25.577 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.577 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-0' for execution.\n"
+       "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"
       ]
      },
      "metadata": {},
@@ -411,11 +411,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\">10:53:25.592 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.592 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-1' for execution.\n"
+       "18:24:37.559 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-8' for execution.\n"
       ]
      },
      "metadata": {},
@@ -424,11 +424,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\">10:53:25.606 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.606 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-8' for execution.\n"
+       "18:24:37.712 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-5' for execution.\n"
       ]
      },
      "metadata": {},
@@ -437,11 +437,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\">10:53:25.613 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.613 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-3' for execution.\n"
+       "18:24:37.731 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-4' for execution.\n"
       ]
      },
      "metadata": {},
@@ -450,11 +450,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\">10:53:25.623 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.623 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-2' for execution.\n"
+       "18:24:37.745 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-7' for execution.\n"
       ]
      },
      "metadata": {},
@@ -463,11 +463,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\">10:53:25.638 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.638 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-4' for execution.\n"
+       "18:24:37.755 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-1' for execution.\n"
       ]
      },
      "metadata": {},
@@ -476,11 +476,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\">10:53:25.713 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.713 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-5' for execution.\n"
+       "18:24:37.771 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-3' for execution.\n"
       ]
      },
      "metadata": {},
@@ -489,11 +489,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\">10:53:25.831 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.831 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-6' for execution.\n"
+       "18:24:37.789 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-0' for execution.\n"
       ]
      },
      "metadata": {},
@@ -502,11 +502,37 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">10:53:25.903 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'masterful-panther'</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: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>\n"
       ],
       "text/plain": [
-       "10:53:25.903 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-7' for execution.\n"
+       "18:24:37.803 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'airborne-chupacabra'\u001b[0m - Submitted task run 'md-2' for execution.\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "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>\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"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "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>\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"
       ]
      },
      "metadata": {},

mlip_arena/tasks/run.py

@@ -240,14 +240,29 @@ def md(
         traj_file.parent.mkdir(parents=True, exist_ok=True)
 
         if restart and traj_file.exists():
-            traj = read(traj_file, index=":")
-            last_atoms = traj[-1]
-            assert isinstance(last_atoms, Atoms)
-            last_step = last_atoms.info.get("step", len(traj) * traj_interval)
-            n_steps -= last_step
-            traj = Trajectory(traj_file, "a", atoms)
-            atoms.set_positions(last_atoms.get_positions())
-            atoms.set_momenta(last_atoms.get_momenta())
+            try:
+                traj = read(traj_file, index=":")
+                last_atoms = traj[-1]
+                assert isinstance(last_atoms, Atoms)
+                last_step = last_atoms.info.get("step", len(traj) * traj_interval)
+                n_steps -= last_step
+                traj = Trajectory(traj_file, "a", atoms)
+                atoms.set_positions(last_atoms.get_positions())
+                atoms.set_momenta(last_atoms.get_momenta())
+            except Exception:
+                traj = Trajectory(traj_file, "w", atoms)
+
+                if not np.isnan(t_schedule).any():
+                    MaxwellBoltzmannDistribution(
+                        atoms=atoms,
+                        temperature_K=t_schedule[last_step],
+                        rng=np.random.default_rng(seed=mb_velocity_seed),
+                    )
+
+                if zero_linear_momentum:
+                    Stationary(atoms)
+                if zero_angular_momentum:
+                    ZeroRotation(atoms)
         else:
             traj = Trajectory(traj_file, "w", atoms)
 

pyproject.toml

@@ -30,6 +30,7 @@ dependencies=[
     "torch",
     "torch_dftd>=0.4.0",
     "huggingface_hub",
+    "datasets",
     "torch-geometric",
     "safetensors",
     "pymatgen"

serve/tasks/homonuclear-diatomics.py

@@ -30,7 +30,7 @@ valid_models = [
 mlip_methods = container.multiselect(
     "MLIPs",
     valid_models,
-    ["EquiformerV2(OC22)", "CHGNet", "M3GNet", "SevenNet", "MACE-MP(M)"],
+    ["EquiformerV2(OC22)", "CHGNet", "M3GNet", "SevenNet", "MACE-MP(M)", "ORB"],
 )
 dft_methods = container.multiselect("DFT Methods", ["GPAW"], [])
 
@@ -171,6 +171,16 @@ for i, symbol in enumerate(chemical_symbols[1:]):
 
     fig.update_layout(
         showlegend=True,
+        legend=dict(
+            orientation="v",
+            x=0.95,
+            xanchor="right",
+            y=1,
+            yanchor="top",
+            bgcolor="rgba(0, 0, 0, 0)"
+            # entrywidth=0.3,
+            # entrywidthmode='fraction',
+        ),
         title_text=f"{name}",
         title_x=0.5,
     )