reform scaffold

.gitignore

@@ -1,5 +1,7 @@
 tests/
 *.out
+*.ipynb
+*.extxyz
 mlip_arena/tasks/*/*/*/
 
 # Byte-compiled / optimized / DLL files

environment.yml

@@ -0,0 +1,335 @@
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - bzip2=1.0.8=h5eee18b_5
+  - ca-certificates=2024.3.11=h06a4308_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - 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
+  - sqlite=3.41.2=h5eee18b_0
+  - tk=8.6.12=h1ccaba5_0
+  - wheel=0.41.2=py311h06a4308_0
+  - xz=5.4.6=h5eee18b_0
+  - zlib=1.2.13=h5eee18b_0
+  - pip:
+      - absl-py==2.1.0
+      - aiofiles==23.2.1
+      - aiohttp==3.9.3
+      - aioitertools==0.11.0
+      - aiosignal==1.3.1
+      - aiosqlite==0.20.0
+      - alembic==1.13.1
+      - alignn==2024.5.27
+      - altair==5.3.0
+      - annotated-types==0.6.0
+      - anyio==3.7.1
+      - appdirs==1.4.4
+      - apprise==1.7.5
+      - ase==3.23.0
+      - asgi-lifespan==2.1.0
+      - asttokens==2.4.1
+      - async-timeout==4.0.3
+      - asyncpg==0.29.0
+      - atomate2==0.0.14.post30+g256b39a1
+      - attrs==23.2.0
+      - autograd==1.5
+      - autoray==0.6.9
+      - bcrypt==4.1.2
+      - bidict==0.23.1
+      - blinker==1.7.0
+      - blosc2==2.7.0
+      - boto3==1.34.74
+      - botocore==1.34.74
+      - cachetools==5.3.3
+      - certifi==2024.7.4
+      - cffi==1.16.0
+      - charset-normalizer==3.3.2
+      - chgnet==0.3.5
+      - click==8.1.7
+      - cloudpickle==3.0.0
+      - colorama==0.4.6
+      - comm==0.2.2
+      - contourpy==1.2.0
+      - coolname==2.2.0
+      - covalent==0.232.0.post1
+      - croniter==2.0.3
+      - cryptography==42.0.5
+      - custodian==2024.3.12
+      - cycler==0.12.1
+      - cython==3.0.10
+      - dask==2024.3.1
+      - dask-jobqueue==0.8.5
+      - dateparser==1.2.0
+      - debugpy==1.8.1
+      - decorator==5.1.1
+      - dgl==2.3.0+cu121
+      - distributed==2024.3.1
+      - dnspython==2.6.1
+      - docker==6.1.3
+      - docker-pycreds==0.4.0
+      - e3nn==0.5.1
+      - email-validator==2.1.1
+      - emmet-core==0.82.1
+      - executing==2.0.1
+      - fairchem-core==1.0.0
+      - fastapi==0.110.0
+      - fastjsonschema==2.20.0
+      - filelock==3.15.4
+      - fireworks==2.0.3
+      - flake8==7.1.0
+      - flask==3.0.2
+      - flask-paginate==2024.3.28
+      - fonttools==4.50.0
+      - frozenlist==1.4.1
+      - fsspec==2024.6.1
+      - furl==2.1.3
+      - future==1.0.0
+      - gitdb==4.0.11
+      - gitpython==3.1.43
+      - google-auth==2.29.0
+      - gpaw==24.7.0b1
+      - greenlet==3.0.3
+      - griffe==0.42.1
+      - grpcio==1.64.1
+      - gunicorn==21.2.0
+      - h11==0.14.0
+      - h2==4.1.0
+      - h5py==3.10.0
+      - hpack==4.0.0
+      - httpcore==1.0.5
+      - httptools==0.6.1
+      - httpx==0.27.0
+      - huggingface-hub==0.22.2
+      - hyperframe==6.0.1
+      - idna==3.7
+      - importlib-metadata==7.1.0
+      - importlib-resources==6.1.3
+      - inflect==7.3.1
+      - ipykernel==6.29.4
+      - ipython==8.22.2
+      - ipywidgets==8.1.2
+      - itsdangerous==2.1.2
+      - jarvis-tools==2024.5.10
+      - jedi==0.19.1
+      - jinja2==3.1.4
+      - jmespath==1.0.1
+      - jobflow==0.1.17
+      - joblib==1.3.2
+      - 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
+      - kiwisolver==1.4.5
+      - kubernetes==29.0.0
+      - latexcodec==3.0.0
+      - lightning-utilities==0.11.2
+      - llvmlite==0.42.0
+      - lmdb==1.4.1
+      - lmdbm==0.0.5
+      - locket==1.0.0
+      - looseversion==1.3.0
+      - mace-torch==0.3.4
+      - maggma==0.64.0
+      - mako==1.3.2
+      - markdown==3.6
+      - markdown-it-py==2.2.0
+      - markupsafe==2.1.5
+      - matgl==1.0.0
+      - 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
+      - mongogrant==0.3.3
+      - mongomock==4.1.2
+      - monty==2024.2.26
+      - 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
+      - natsort==8.4.0
+      - nbformat==5.10.4
+      - ndindex==1.8
+      - nest-asyncio==1.6.0
+      - networkx==3.3
+      - numba==0.59.1
+      - numexpr==2.10.1
+      - numpy==1.26.4
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==8.9.2.26
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - 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-nvtx-cu12==12.1.105
+      - oauthlib==3.2.2
+      - opt-einsum==3.3.0
+      - opt-einsum-fx==0.1.4
+      - orderedmultidict==1.0.1
+      - orjson==3.10.0
+      - packaging==24.0
+      - palettable==3.3.3
+      - pandas==2.2.2
+      - paramiko==3.4.0
+      - parso==0.8.3
+      - partd==1.4.1
+      - pathspec==0.12.1
+      - pendulum==2.1.2
+      - pennylane==0.32.0
+      - pennylane-lightning==0.33.1
+      - pexpect==4.9.0
+      - pillow==10.2.0
+      - 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
+      - prompt-toolkit==3.0.43
+      - protobuf==4.25.3
+      - psutil==6.0.0
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.2
+      - py-cpuinfo==9.0.0
+      - pyarrow==16.1.0
+      - pyasn1==0.6.0
+      - pyasn1-modules==0.4.0
+      - pybtex==0.24.0
+      - pycodestyle==2.12.0
+      - pycparser==2.22
+      - pydantic==2.6.4
+      - pydantic-core==2.16.3
+      - pydantic-settings==2.2.1
+      - pydash==8.0.0
+      - pydeck==0.9.1
+      - pydocstyle==6.3.0
+      - pyflakes==3.2.0
+      - pygments==2.17.2
+      - pymatgen==2024.4.13
+      - pymongo==4.6.3
+      - pynacl==1.5.0
+      - 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-multipart==0.0.9
+      - python-slugify==8.0.4
+      - python-socketio==5.11.2
+      - pytorch-lightning==2.2.1
+      - pytz==2024.1
+      - pytzdata==2020.1
+      - pyyaml==6.0.1
+      - pyzmq==25.1.2
+      - 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
+      - rich==13.3.5
+      - rpds-py==0.18.0
+      - rsa==4.9
+      - ruamel-yaml==0.17.40
+      - ruamel-yaml-clib==0.2.8
+      - rustworkx==0.14.2
+      - s3transfer==0.10.1
+      - safetensors==0.4.2
+      - scikit-learn==1.4.1.post1
+      - scipy==1.14.0
+      - semantic-version==2.10.0
+      - sentinels==1.0.0
+      - sentry-sdk==2.7.1
+      - setproctitle==1.3.3
+      - shellingham==1.5.4
+      - simple-websocket==1.0.0
+      - simplejson==3.19.2
+      - six==1.16.0
+      - smart-open==7.0.4
+      - smmap==5.0.1
+      - sniffio==1.3.1
+      - snowballstemmer==2.2.0
+      - sortedcontainers==2.4.0
+      - spglib==2.3.1
+      - sqlalchemy==1.4.52
+      - sqlalchemy-utils==0.41.2
+      - sshtunnel==0.4.0
+      - stack-data==0.6.3
+      - starlette==0.36.3
+      - streamlit==1.36.0
+      - submitit==1.5.1
+      - sympy==1.12.1
+      - 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
+      - text-unidecode==1.3
+      - threadpoolctl==3.4.0
+      - toml==0.10.2
+      - toolz==0.12.1
+      - torch==2.2.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
+      - torchmetrics==1.3.2
+      - tornado==6.4
+      - tqdm==4.66.4
+      - trainstation==1.0
+      - traitlets==5.14.2
+      - triton==2.2.0
+      - typeguard==4.3.0
+      - typer==0.12.0
+      - typer-cli==0.12.0
+      - typer-slim==0.12.0
+      - typing-extensions==4.12.2
+      - tzdata==2024.1
+      - tzlocal==5.2
+      - ujson==5.9.0
+      - uncertainties==3.1.7
+      - urllib3==2.2.2
+      - uvicorn==0.18.3
+      - uvloop==0.19.0
+      - wandb==0.17.4
+      - watchdog==4.0.0
+      - watchfiles==0.21.0
+      - wcwidth==0.2.13
+      - websocket-client==1.7.0
+      - websockets==12.0
+      - werkzeug==3.0.1
+      - widgetsnbextension==4.0.10
+      - wrapt==1.16.0
+      - wsproto==1.2.0
+      - xmltodict==0.13.0
+      - yarl==1.9.4
+      - zict==3.0.0
+      - zipp==3.18.1

mlip_arena/models/chgnet.py

@@ -0,0 +1,60 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+from ase import Atoms
+from ase.calculators.calculator import all_changes
+from huggingface_hub import hf_hub_download
+from torch_geometric.data import Data
+
+from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP
+
+
+class CHGNetCalculator(MLIPCalculator):
+    def __init__(
+        self,
+        device: torch.device | None = None,
+        restart=None,
+        atoms=None,
+        directory=".",
+        **kwargs,
+    ):
+        super().__init__(restart=restart, atoms=atoms, directory=directory, **kwargs)
+
+        self.name: str = self.__class__.__name__
+
+        fpath = hf_hub_download(
+            repo_id="cyrusyc/mace-universal",
+            subfolder="pretrained",
+            filename="2023-12-12-mace-128-L1_epoch-199.model",
+            revision="main",
+        )
+
+        self.device = device or torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        self.model = torch.load(fpath, map_location=self.device)
+
+        self.implemented_properties = ["energy", "forces", "stress"]
+
+    def calculate(
+        self, atoms: Atoms, properties: list[str], system_changes: list = all_changes
+    ):
+        """Calculate energies and forces for the given Atoms object"""
+        super().calculate(atoms, properties, system_changes)
+
+        output = self.forward(atoms)
+
+        self.results = {}
+        if "energy" in properties:
+            self.results["energy"] = output["energy"].item()
+        if "forces" in properties:
+            self.results["forces"] = output["forces"].cpu().detach().numpy()
+        if "stress" in properties:
+            self.results["stress"] = output["stress"].cpu().detach().numpy()
+
+    def forward(self, x: Data | Atoms) -> dict[str, torch.Tensor]:
+        """Implement data conversion, graph creation, and model forward pass"""
+        # TODO
+        raise NotImplementedError

mlip_arena/tasks/diatomics/mlip.py

@@ -0,0 +1,165 @@
+from datetime import timedelta
+from typing import Union
+
+# import covalent as ct
+import numpy as np
+import pandas as pd
+import torch
+from ase import Atoms
+from ase.calculators.calculator import Calculator
+from ase.data import chemical_symbols
+from dask.distributed import Client
+from dask_jobqueue import SLURMCluster
+from prefect import flow, task
+from prefect.tasks import task_input_hash
+from prefect_dask import DaskTaskRunner
+
+from mlip_arena.models import MLIPCalculator
+from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
+
+cluster_kwargs = {
+    "cores": 4,
+    "memory": "64 GB",
+    "shebang": "#!/bin/bash",
+    "account": "m3828",
+    "walltime": "00:10:00",
+    "job_mem": "0",
+    "job_script_prologue": ["source ~/.bashrc"],
+    "job_directives_skip": ["-n", "--cpus-per-task"],
+    "job_extra_directives": ["-q debug", "-C gpu"],
+}
+
+cluster = SLURMCluster(**cluster_kwargs)
+cluster.scale(jobs=10)
+client = Client(cluster)
+
+
+@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=1))
+def calculate_single_diatomic(
+    calculator_name: str | EXTMLIPEnum,
+    calculator_kwargs: dict | None,
+    atom1: str,
+    atom2: str,
+    rmin: float = 0.1,
+    rmax: float = 6.5,
+    npts: int = int(1e3),
+):
+
+    calculator_kwargs = calculator_kwargs or {}
+
+    if isinstance(calculator_name, EXTMLIPEnum) and calculator_name in EXTMLIPEnum:
+        calc = external_ase_calculator(calculator_name, **calculator_kwargs)
+    elif calculator_name in MLIPMap:
+        calc = MLIPMap[calculator_name](**calculator_kwargs)
+
+    a = 2 * rmax
+
+    rs = np.linspace(rmin, rmax, npts)
+    e = np.zeros_like(rs)
+    f = np.zeros_like(rs)
+
+    da = atom1 + atom2
+
+    for i, r in enumerate(rs):
+
+        positions = [
+            [0, 0, 0],
+            [r, 0, 0],
+        ]
+
+        # Create the unit cell with two atoms
+        atoms = Atoms(da, positions=positions, cell=[a, a, a])
+
+        atoms.calc = calc
+
+        e[i] = atoms.get_potential_energy()
+        f[i] = np.inner(np.array([1, 0, 0]), atoms.get_forces()[1])
+
+    return {"r": rs, "E": e, "F": f, "da": da}
+
+
+@flow
+def calculate_multiple_diatomics(calculator_name, calculator_kwargs):
+
+    futures = []
+    for symbol in chemical_symbols:
+        if symbol == "X":
+            continue
+        futures.append(
+            calculate_single_diatomic.submit(
+                calculator_name, calculator_kwargs, symbol, symbol
+            )
+        )
+
+    return [i for future in futures for i in future.result()]
+
+
+@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True)
+def calculate_homonuclear_diatomics(calculator_name, calculator_kwargs):
+
+    curves = calculate_multiple_diatomics(calculator_name, calculator_kwargs)
+
+    pd.DataFrame(curves).to_csv(f"homonuclear-diatomics-{calculator_name}.csv")
+
+
+# with plt.style.context("default"):
+
+#     SMALL_SIZE = 6
+#     MEDIUM_SIZE = 8
+#     LARGE_SIZE = 10
+
+#     LINE_WIDTH = 1
+
+#     plt.rcParams.update(
+#         {
+#             "pgf.texsystem": "pdflatex",
+#             "font.family": "sans-serif",
+#             "text.usetex": True,
+#             "pgf.rcfonts": True,
+#             "figure.constrained_layout.use": True,
+#             "axes.labelsize": MEDIUM_SIZE,
+#             "axes.titlesize": MEDIUM_SIZE,
+#             "legend.frameon": False,
+#             "legend.fontsize": MEDIUM_SIZE,
+#             "legend.loc": "best",
+#             "lines.linewidth": LINE_WIDTH,
+#             "xtick.labelsize": SMALL_SIZE,
+#             "ytick.labelsize": SMALL_SIZE,
+#         }
+#     )
+
+#     fig, ax = plt.subplots(layout="constrained", figsize=(3, 2), dpi=300)
+
+#     color = "tab:red"
+#     ax.plot(rs, e, color=color, zorder=1)
+
+#     ax.axhline(ls="--", color=color, alpha=0.5, lw=0.5 * LINE_WIDTH)
+
+#     ylo, yhi = ax.get_ylim()
+#     ax.set(xlabel=r"r [$\AA]$", ylim=(max(-7, ylo), min(5, yhi)))
+#     ax.set_ylabel(ylabel="E [eV]", color=color)
+#     ax.tick_params(axis="y", labelcolor=color)
+#     ax.text(0.8, 0.85, da, fontsize=LARGE_SIZE, transform=ax.transAxes)
+
+#     color = "tab:blue"
+
+#     at = ax.twinx()
+#     at.plot(rs, f, color=color, zorder=0, lw=0.5 * LINE_WIDTH)
+
+#     at.axhline(ls="--", color=color, alpha=0.5, lw=0.5 * LINE_WIDTH)
+
+#     ylo, yhi = at.get_ylim()
+#     at.set(
+#         xlabel=r"r [$\AA]$",
+#         ylim=(max(-20, ylo), min(20, yhi)),
+#     )
+#     at.set_ylabel(ylabel="F [eV/$\AA$]", color=color)
+#     at.tick_params(axis="y", labelcolor=color)
+
+#     plt.show()
+
+
+if __name__ == "__main__":
+    calculate_homonuclear_diatomics(
+        EXTMLIPEnum.MACE, dict(model="medium", device="cuda")
+    )

mlip_arena/tasks/diatomics/vasp/run.ipynb

@@ -1,509 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "from datetime import timedelta\n",
-    "from typing import Union\n",
-    "import os\n",
-    "from pathlib import Path\n",
-    "from itertools import combinations, combinations_with_replacement\n",
-    "import psutil\n",
-    "import subprocess\n",
-    "\n",
-    "import numpy as np\n",
-    "import pandas as pd\n",
-    "import torch\n",
-    "from ase import Atoms, Atom\n",
-    "from ase.calculators.calculator import Calculator\n",
-    "from ase.calculators.singlepoint import SinglePointCalculator\n",
-    "from ase.calculators.vasp import Vasp\n",
-    "from ase.data import chemical_symbols, covalent_radii, vdw_alvarez\n",
-    "from ase.io import read, write\n",
-    "from dask.distributed import Client\n",
-    "from dask_jobqueue import SLURMCluster\n",
-    "from prefect import flow, task\n",
-    "from prefect.tasks import task_input_hash\n",
-    "from prefect_dask import DaskTaskRunner\n",
-    "\n",
-    "from pymatgen.core import Element\n",
-    "from pymatgen.io.ase import AseAtomsAdaptor\n",
-    "from pymatgen.io.vasp.inputs import Kpoints\n",
-    "from pymatgen.command_line.chargemol_caller import ChargemolAnalysis\n",
-    "\n",
-    "from mlip_arena.models import MLIPCalculator\n",
-    "from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator\n",
-    "from jobflow import run_locally\n",
-    "\n",
-    "from atomate2.vasp.jobs.mp import MPGGAStaticMaker\n",
-    "from atomate2.vasp.sets.mp import MPGGAStaticSetGenerator"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": true,
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "\n",
-    "nodes_per_alloc = 1\n",
-    "cpus_per_task = 16\n",
-    "gpus_per_node = 1\n",
-    "# tasks_per_node = int(128/4)\n",
-    "ntasks = 1\n",
-    "\n",
-    "\n",
-    "cluster_kwargs = dict(\n",
-    "    cores=1,\n",
-    "    memory=\"64 GB\",\n",
-    "    shebang=\"#!/bin/bash\",\n",
-    "    account=\"m3828\",\n",
-    "    walltime=\"01:00:00\",\n",
-    "    # processes=16,\n",
-    "    # nanny=True,\n",
-    "    job_mem=\"0\",\n",
-    "    job_script_prologue=[\n",
-    "        \"source ~/.bashrc\",\n",
-    "        \"module load python\",\n",
-    "        \"source activate /pscratch/sd/c/cyrisinstance.conda/mlip-arena\",\n",
-    "        \"module load vasp/6.4.1-gpu\",\n",
-    "        \n",
-    "        \"export DDEC6_ATOMIC_DENSITIES_DIR='/global/homes/c/cyrusyc/chargemol/atomic_densities/'\",\n",
-    "        f\"export OMP_NUM_THREADS={gpus_per_node}\",\n",
-    "        \"export OMP_PLACES=threads\",\n",
-    "        \"export OMP_PROC_BIND=spread\",\n",
-    "        f\"export ASE_VASP_COMMAND='srun -n {ntasks} -c {4*gpus_per_node} --cpu-bind=cores --gpus-per-node {gpus_per_node} vasp_ncl'\"\n",
-    "        # f\"export ASE_VASP_COMMAND='srun -N {nodes_per_alloc} --ntasks-per-node={tasks_per_node} -c {cpus_per_task} --cpu-bind=cores vasp_std'\"\n",
-    "        # \"export ATOMATE2_CONFIG_FILE='/global/homes/c/cyrusyc/atomate2/config/atomate2-prefect-cpu-node.yaml'\"\n",
-    "    ],\n",
-    "    job_directives_skip=[\"-n\", \"--cpus-per-task\", \"-J\"],\n",
-    "    job_extra_directives=[\n",
-    "        \"-J diatomics\",\n",
-    "        \"-q regular\",\n",
-    "        f\"-N {nodes_per_alloc}\",\n",
-    "        \"-C gpu\",\n",
-    "        # \"-n 1\",\n",
-    "        # \"-c 16\",\n",
-    "        # \"--gpus-per-task=1\",\n",
-    "        # \"--threads-per-task=1\",\n",
-    "        # \"--gpu-bind=single:1\",\n",
-    "        # \"--comment=00:20:00\",\n",
-    "        # \"--time-min=00:05:00\",\n",
-    "        # \"--signal=B:USR1@60\",\n",
-    "        # \"--requeue\",\n",
-    "        # \"--open-mode=append\",\n",
-    "        # \"--mail-type=end,requeue\",\n",
-    "        # \"--mail-user=cyrusyc@lbl.gov\",\n",
-    "    ],\n",
-    "    # python=\"srun python\",\n",
-    "    death_timeout=86400, #float('inf')\n",
-    ")\n",
-    "\n",
-    "\n",
-    "cluster = SLURMCluster(**cluster_kwargs)\n",
-    "print(cluster.job_script())\n",
-    "# cluster.scale(3)\n",
-    "cluster.adapt(minimum_jobs=50, maximum_jobs=100)\n",
-    "client = Client(cluster)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "\n",
-    "\n",
-    "@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=24), log_prints=True)\n",
-    "def calculate_single_diatomic(\n",
-    "    calculator: str | EXTMLIPEnum | Calculator,\n",
-    "    calculator_kwargs: dict | None,\n",
-    "    atom1: str,\n",
-    "    atom2: str,\n",
-    "    rmin: float = 1.25,\n",
-    "    rmax: float = 6.25,\n",
-    "    rstep: float = 0.2,\n",
-    "    magnetism: str = \"FM\"\n",
-    "):\n",
-    "\n",
-    "    calculator_kwargs = calculator_kwargs or {}\n",
-    "\n",
-    "    if isinstance(calculator, str) and calculator.lower() == 'vasp-mp-gga':\n",
-    "        calc = Vasp(**calculator_kwargs)\n",
-    "        calc.name = 'vasp-mp-gga'\n",
-    "        # calc.name='atomate2'\n",
-    "    elif isinstance(calculator, EXTMLIPEnum) and calculator in EXTMLIPEnum:\n",
-    "        calc = external_ase_calculator(calculator, **calculator_kwargs)\n",
-    "    elif calculator in MLIPMap:\n",
-    "        calc = MLIPMap[calculator](**calculator_kwargs)\n",
-    "    elif issubclass(calculator, Calculator):\n",
-    "        calc = calculator(**calculator_kwargs)\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",
-    "    f = np.zeros_like(rs)\n",
-    "\n",
-    "    da = atom1 + atom2\n",
-    "    \n",
-    "    assert isinstance(calc, Calculator)\n",
-    "    \n",
-    "    out_dir = Path(str(da + f\"_{magnetism}\"))\n",
-    "    os.makedirs(out_dir, exist_ok=True)\n",
-    "    \n",
-    "    calc.directory = out_dir\n",
-    "    \n",
-    "    print(f\"write output to {calc.directory}\")\n",
-    "    \n",
-    "    element = Element(atom1)\n",
-    "        \n",
-    "    try:\n",
-    "        m = element.valence[1]\n",
-    "        if element.valence == (0, 2):\n",
-    "            m = 0\n",
-    "    except:\n",
-    "        m = 0\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",
-    "    if magnetism == 'FM':\n",
-    "        if m == 0:\n",
-    "            return {}\n",
-    "        magmoms = [m, m]\n",
-    "    elif magnetism == 'AFM':\n",
-    "        if m == 0:\n",
-    "            return {}\n",
-    "        magmoms = [m, -m]\n",
-    "    elif magnetism == 'NM':\n",
-    "        magmoms = [0, 0]\n",
-    "    \n",
-    "    traj_fpath = out_dir / \"traj.extxyz\"\n",
-    "    \n",
-    "    skip = 0\n",
-    "    if traj_fpath.exists():\n",
-    "        traj = read(traj_fpath, index=\":\")\n",
-    "        skip = len(traj)\n",
-    "        atoms = traj[-1]\n",
-    "    else:\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",
-    "    # \n",
-    "    \n",
-    "    structure = AseAtomsAdaptor().get_structure(atoms)\n",
-    "            \n",
-    "    if magnetism == 'FM':\n",
-    "        I_CONSTRAINED_M = 2\n",
-    "        LAMBDA = 10\n",
-    "        M_CONSTR = [0, 0, 1, 0, 0, 1] # \" \".join(map(str, [0, 0, 1])) + \" \" + \" \".join(map(str, [0, 0, 1]))\n",
-    "    elif magnetism == 'AFM':\n",
-    "        I_CONSTRAINED_M = 2\n",
-    "        LAMBDA = 10\n",
-    "        M_CONSTR = [0, 0, 1, 0, 0, -1] # \" \".join(map(str, [0, 0, 1])) + \" \" + \" \".join(map(str, [0, 0, -1]))\n",
-    "    elif magnetism == 'NM':\n",
-    "        I_CONSTRAINED_M = 1\n",
-    "        LAMBDA = 10\n",
-    "        M_CONSTR = [0, 0, 0, 0, 0, 0] #\" \".join(map(str, [0, 0, 0])) + \" \" + \" \".join(map(str, [0, 0, 0]))\n",
-    "\n",
-    "    input_set_generator = MPGGAStaticSetGenerator(\n",
-    "        user_incar_settings=dict(\n",
-    "            ISYM   = 0, # symmetry is off\n",
-    "            ISPIN  = 2,\n",
-    "            ISMEAR = 0,    # Gaussian smearing, otherwise negative occupancies might come up\n",
-    "            SIGMA  = 0.002, # tiny smearing width to safely break symmetry\n",
-    "            AMIX   = 0.2,    # mixing set manually\n",
-    "            BMIX   = 0.0001,\n",
-    "            LSUBROT= True, # spin orbit coupling (non collinear)\n",
-    "            ALGO = \"Accurate\",\n",
-    "            PREC = \"High\",\n",
-    "            ENCUT = 1000,\n",
-    "            ENAUG = 2000,\n",
-    "            ISTART = 1,\n",
-    "            ICHARG = 1,\n",
-    "            NELM   = 200,\n",
-    "            TIME = 0.2,\n",
-    "            LELF   = False,\n",
-    "            LMAXMIX=max(max(map(lambda a: \"spdf\".index(Element(a.symbol).block) * 2, atoms)), 2),\n",
-    "            LMIXTAU=False,\n",
-    "            VOSKOWN = 1,\n",
-    "            I_CONSTRAINED_M = I_CONSTRAINED_M,\n",
-    "            M_CONSTR = M_CONSTR,\n",
-    "            LAMBDA = LAMBDA,\n",
-    "            # performance\n",
-    "            # lplane=False,\n",
-    "            # npar=int(sqrt(ncpus)),\n",
-    "            # nsim=1,\n",
-    "            # LPLANE = True,\n",
-    "            # # NCORE = 128,\n",
-    "            # LSCALU = False,\n",
-    "            # NSIM = 4,\n",
-    "            # LPLANE = False,\n",
-    "            # NPAR  = 16,\n",
-    "            # NSIM   = 1,\n",
-    "            # LSCALU = False,\n",
-    "            # GPU\n",
-    "            KPAR = gpus_per_node,\n",
-    "            NSIM = 64,\n",
-    "\n",
-    "            LVTOT  = False,\n",
-    "            LAECHG = True,  # AECCARs\n",
-    "            LASPH  = True,   # aspherical charge density\n",
-    "            LCHARG = True,  # CHGCAR\n",
-    "            LWAVE  = True\n",
-    "        ),\n",
-    "        user_kpoints_settings=Kpoints(), # Gamma point only\n",
-    "        user_potcar_settings={\n",
-    "            \"Yb\": \"Yb_3\"\n",
-    "        },\n",
-    "        sort_structure=False\n",
-    "    )\n",
-    "\n",
-    "    vis = input_set_generator.get_input_set(structure=structure)\n",
-    "    vis.incar.pop(\"MAGMOM\")\n",
-    "\n",
-    "    incar = {key.lower(): value for key, value in vis.incar.items()} \n",
-    "    calc.set(kpts=1, gamma=True, **incar)\n",
-    "        \n",
-    "    atoms.calc = calc\n",
-    "\n",
-    "    for i, r in enumerate(np.flip(rs)):\n",
-    "\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",
-    "        if i > 0:   \n",
-    "            magmoms = atoms.get_magnetic_moments()\n",
-    "        \n",
-    "        atoms.set_initial_magnetic_moments(magmoms)\n",
-    "        atoms.set_positions(positions)\n",
-    "\n",
-    "        print(f\"{atoms} separated by {r} A ({i+1}/{len(rs)})\")\n",
-    "        \n",
-    "\n",
-    "        e[i] = atoms.get_potential_energy()\n",
-    "        f[i] = np.inner(np.array([1, 0, 0]), atoms.get_forces()[1])\n",
-    "        \n",
-    "        atoms.calc.results.update(dict(\n",
-    "            magmoms=atoms.get_magnetic_moments()\n",
-    "        ))\n",
-    "        \n",
-    "        write(out_dir / \"traj.extxyz\", atoms, append=\"a\")\n",
-    "        \n",
-    "#         additional_results = {}\n",
-    "        \n",
-    "#         try:\n",
-    "#             ncpus = psutil.cpu_count(logical=True)\n",
-    "#             nthreads = os.environ[\"OMP_NUM_THREADS\"]\n",
-    "#             subprocess.run([\"export\", f\"OMP_NUM_THREADS={ncpus}\"], shell=True)\n",
-    "            \n",
-    "#             ca = ChargemolAnalysis(path=out_dir)\n",
-    "        \n",
-    "#             if charges := ca.ddec_charges:\n",
-    "#                 additional_results[\"charges\"] = np.array(charges)\n",
-    "#             if dipoles := ca.dipoles:\n",
-    "#                 additional_results[\"dipoles\"] = np.array(dipoles)\n",
-    "#             if magmoms := ca.ddec_spin_moments:\n",
-    "#                 additional_results[\"magmoms\"] = np.array(magmoms)\n",
-    "            \n",
-    "#             subprocess.run([\"export\", f\"OMP_NUM_THREADS={nthreads}\"], shell=True)\n",
-    "#         except:\n",
-    "#             print(\"DDEC failed\")\n",
-    "            \n",
-    "        \n",
-    "#         atoms.calc.results.update(additional_results)\n",
-    "        \n",
-    "    return {\"r\": rs, \"E\": e, \"F\": f, \"da\": da}\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True)\n",
-    "def calculate_multiple_diatomics(calculator_name, calculator_kwargs):\n",
-    "\n",
-    "    futures = []\n",
-    "    for sa in chemical_symbols:\n",
-    "        \n",
-    "        s = set([sa])\n",
-    "        \n",
-    "        if 'X' in s:\n",
-    "            continue\n",
-    "        \n",
-    "        atom = Atom(sa)\n",
-    "        rmin = covalent_radii[atom.number] * 2 * 0.6\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.2 #if rmin < 1 else 0.4\n",
-    "        \n",
-    "        futures.append(\n",
-    "            calculate_single_diatomic.submit(\n",
-    "                calculator_name, calculator_kwargs, sa, sa,\n",
-    "                rmin=rmin, rmax=rmax,\n",
-    "                rstep=rstep,\n",
-    "                magnetism=\"FM\"\n",
-    "                # npts=16 if 'H' in s else 21\n",
-    "            )\n",
-    "        )\n",
-    "        futures.append(\n",
-    "            calculate_single_diatomic.submit(\n",
-    "                calculator_name, calculator_kwargs, sa, sa,\n",
-    "                rmin=rmin, rmax=rmax,\n",
-    "                rstep=rstep, #0.1 if rmin < 1 else 0.25,\n",
-    "                magnetism=\"AFM\"\n",
-    "                # npts=16 if 'H' in s else 21\n",
-    "            )\n",
-    "        )\n",
-    "        futures.append(\n",
-    "            calculate_single_diatomic.submit(\n",
-    "                calculator_name, calculator_kwargs, sa, sa,\n",
-    "                rmin=rmin, rmax=rmax,\n",
-    "                rstep=rstep, #0.1 if rmin < 1 else 0.25,\n",
-    "                magnetism=\"NM\"\n",
-    "                # npts=16 if 'H' in s else 21\n",
-    "            )\n",
-    "        )\n",
-    "#     for sa, sb in combinations_with_replacement(chemical_symbols, 2):\n",
-    "        \n",
-    "#         if 'X' in set([sa, sb]):\n",
-    "#             continue\n",
-    "        \n",
-    "#         futures.append(\n",
-    "#             calculate_single_diatomic.submit(\n",
-    "#                 calculator_name, calculator_kwargs, sa, sb,\n",
-    "#                 rmin=0.5, rmax=6.5,\n",
-    "#                 npts=16\n",
-    "#             )\n",
-    "#         )\n",
-    "\n",
-    "    return [i for future in futures for i in future.result()]\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": true,
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "calculate_multiple_diatomics(\n",
-    "    \"vasp-mp-gga\", \n",
-    "    dict(\n",
-    "        xc=\"pbe\",\n",
-    "        kpts=1,\n",
-    "        # Massively parallel machines (Cray)\n",
-    "        # lplane=False,\n",
-    "        # npar=int(sqrt(ncpus)),\n",
-    "        # nsim=1\n",
-    "        # Multicore modern linux machines\n",
-    "        # lplane=True,\n",
-    "        # npar=2,\n",
-    "        # lscalu=False,\n",
-    "        # nsim=4\n",
-    "    )\n",
-    ")\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": true,
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "\n",
-    "calculate_homonuclear_diatomics(\n",
-    "    \"vasp-mp-gga\", \n",
-    "    dict(\n",
-    "        xc=\"pbe\",\n",
-    "        kpts=1,\n",
-    "        # Massively parallel machines (Cray)\n",
-    "        # lplane=False,\n",
-    "        # npar=int(sqrt(ncpus)),\n",
-    "        # nsim=1\n",
-    "        # Multicore modern linux machines\n",
-    "        # lplane=True,\n",
-    "        # npar=2,\n",
-    "        # lscalu=False,\n",
-    "        # nsim=4\n",
-    "    )\n",
-    ")\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": 4
-}

mlip_arena/tasks/stability/run.py

@@ -0,0 +1,3 @@
+
+from mlip_arena.tasks.utils import _valid_dynamics, _preset_dynamics
+

mlip_arena/tasks/utils.py

@@ -0,0 +1,161 @@
+import os, glob
+from pathlib import Path
+from ase.io import read, write
+from ase import units
+from ase import Atoms, units
+from ase.calculators.calculator import Calculator
+from ase.data import chemical_symbols
+from ase.md.andersen import Andersen
+from ase.md.langevin import Langevin
+from ase.md.md import MolecularDynamics
+from ase.md.npt import NPT
+from ase.md.nptberendsen import NPTBerendsen
+from ase.md.nvtberendsen import NVTBerendsen
+from ase.md.velocitydistribution import (
+    MaxwellBoltzmannDistribution,
+    Stationary,
+    ZeroRotation,
+)
+from ase.md.verlet import VelocityVerlet
+from dask.distributed import Client
+from dask_jobqueue import SLURMCluster
+from jobflow import Maker
+from prefect import flow, task
+from prefect.tasks import task_input_hash
+from prefect_dask import DaskTaskRunner
+from pymatgen.io.ase import AseAtomsAdaptor
+from scipy.interpolate import interp1d
+from scipy.linalg import schur
+
+from mlip_arena.models import MLIPCalculator
+from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
+from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+from mp_api.client import MPRester
+
+from fireworks import LaunchPad
+from atomate2.vasp.flows.core import RelaxBandStructureMaker
+from atomate2.vasp.flows.mp import MPGGADoubleRelaxStaticMaker
+from atomate2.vasp.powerups import add_metadata_to_flow
+from atomate2.forcefields.md import (
+    CHGNetMDMaker,
+    GAPMDMaker,
+    M3GNetMDMaker,
+    MACEMDMaker,
+    NequipMDMaker,
+)
+from atomate2.forcefields.utils import MLFF
+from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.transformations.advanced_transformations import CubicSupercellTransformation
+from jobflow.managers.fireworks import flow_to_workflow
+from jobflow import run_locally, SETTINGS
+from tqdm.auto import tqdm
+
+from datetime import timedelta, datetime
+from typing import Literal, Sequence, Tuple
+
+import numpy as np
+import torch
+from pymatgen.core.structure import Structure
+
+from ase.calculators.mixing import SumCalculator
+from scipy.interpolate import interp1d
+
+from ase.io.trajectory import Trajectory
+
+
+_valid_dynamics: dict[str, tuple[str, ...]] = {
+    "nve": ("velocityverlet",),
+    "nvt": ("nose-hoover", "langevin", "andersen", "berendsen"),
+    "npt": ("nose-hoover", "berendsen"),
+}
+
+_preset_dynamics: dict = {
+    "nve_velocityverlet": VelocityVerlet,
+    "nvt_andersen": Andersen,
+    "nvt_berendsen": NVTBerendsen,
+    "nvt_langevin": Langevin,
+    "nvt_nose-hoover": NPT,
+    "npt_berendsen": NPTBerendsen,
+    "npt_nose-hoover": NPT,
+}
+
+def _interpolate_quantity(values: Sequence | np.ndarray, n_pts: int) -> np.ndarray:
+    """Interpolate temperature / pressure on a schedule."""
+    n_vals = len(values)
+    return np.interp(
+        np.linspace(0, n_vals - 1, n_pts + 1),
+        np.linspace(0, n_vals - 1, n_vals),
+        values,
+    )
+
+def _get_ensemble_schedule(
+        ensemble: Literal["nve", "nvt", "npt"] = "nvt",
+        n_steps: int = 1000,
+        temperature: float | Sequence | np.ndarray | None = 300.0,
+        pressure: float | Sequence | np.ndarray | None = None
+        ) -> Tuple[np.ndarray, np.ndarray]:
+    if ensemble == "nve":
+        # Disable thermostat and barostat
+        temperature = np.nan
+        pressure = np.nan
+        t_schedule = np.full(n_steps + 1, temperature)
+        p_schedule = np.full(n_steps + 1, pressure)
+        return t_schedule, p_schedule
+
+    if isinstance(temperature, Sequence) or (
+        isinstance(temperature, np.ndarray) and temperature.ndim == 1
+    ):
+        t_schedule = _interpolate_quantity(temperature, n_steps)
+    # NOTE: In ASE Langevin dynamics, the temperature are normally
+    # scalars, but in principle one quantity per atom could be specified by giving
+    # an array. This is not implemented yet here.
+    else:
+        t_schedule = np.full(n_steps + 1, temperature)
+
+    if ensemble == "nvt":
+        pressure = np.nan
+        p_schedule = np.full(n_steps + 1, pressure)
+        return t_schedule, p_schedule
+
+    if isinstance(pressure, Sequence) or (
+        isinstance(pressure, np.ndarray) and pressure.ndim == 1
+    ):
+        p_schedule = _interpolate_quantity(pressure, n_steps)
+    elif isinstance(pressure, np.ndarray) and pressure.ndim == 4:
+        p_schedule = interp1d(
+            np.arange(n_steps + 1), pressure, kind="linear"
+        )
+        assert isinstance(p_schedule, np.ndarray)
+    else:
+        p_schedule = np.full(n_steps + 1, pressure)
+    
+    return t_schedule, p_schedule
+
+def _get_ensemble_defaults(
+        ensemble: Literal["nve", "nvt", "npt"],
+        dynamics: str | MolecularDynamics,
+        t_schedule: np.ndarray,
+        p_schedule: np.ndarray,
+        ase_md_kwargs: dict | None = None) -> dict:
+    """Update ASE MD kwargs"""
+    ase_md_kwargs = ase_md_kwargs or {}
+
+    if ensemble == "nve":
+        ase_md_kwargs.pop("temperature", None)
+        ase_md_kwargs.pop("temperature_K", None)
+        ase_md_kwargs.pop("externalstress", None)
+    elif ensemble == "nvt":
+        ase_md_kwargs["temperature_K"] = t_schedule[0]
+        ase_md_kwargs.pop("externalstress", None)
+    elif ensemble == "npt":
+        ase_md_kwargs["temperature_K"] = t_schedule[0]
+        ase_md_kwargs["externalstress"] = p_schedule[0] * 1e3 * units.bar
+
+    if isinstance(dynamics, str) and dynamics.lower() == "langevin":
+        ase_md_kwargs["friction"] = ase_md_kwargs.get(
+            "friction",
+            10.0 * 1e-3 / units.fs,  # Same default as in VASP: 10 ps^-1
+        )
+        
+    return ase_md_kwargs
+

pyproject.toml

@@ -46,6 +46,7 @@ Issues = "https://github.com/atomind-ai/mlip-arena/issues"
 
 [tool.ruff]
 # Exclude a variety of commonly ignored directories.
+extend-include = ["*.ipynb"]
 exclude = [
     ".bzr",
     ".direnv",

scripts/install-pyg.sh

@@ -0,0 +1,15 @@
+
+
+# PyTorch Geometric (OCP)
+TORCH=2.2.0
+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

serve/models/alerts.py

@@ -0,0 +1,4 @@
+import streamlit as st
+
+
+st.markdown("# Alerts")

serve/models/bugs.py

@@ -0,0 +1,4 @@
+import streamlit as st
+
+
+st.markdown("# Bugs")

serve/tasks/homonuclear-diatomics.py

@@ -136,7 +136,7 @@ for i, symbol in enumerate(chemical_symbols[1:]):
     )
 
     # Set x-axis title
-    fig.update_xaxes(title_text="Bond length (Å)")
+    fig.update_xaxes(title_text="Bond length [Å]")
 
     # Set y-axes titles
     if energy_plot: