add notebooks

mlip_arena/tasks/diatomics/alignn/run.ipynb

@@ -0,0 +1,243 @@
+{
+ "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
+}

mlip_arena/tasks/diatomics/gpaw/run.ipynb

@@ -0,0 +1,281 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "08ee4c09-7fb5-4ce1-a7a3-5c5c52a6d4d5",
+   "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",
+    "from gpaw import GPAW, PW, FermiDirac, LCAO\n",
+    "from gpaw import Davidson\n",
+    "from gpaw import Mixer, MixerSum, MixerDif\n",
+    "from gpaw.directmin.etdm_lcao import LCAOETDM\n",
+    "from gpaw.cdft.cdft import CDFT\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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90887faa-1601-4c4c-9c44-d16731471d7f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "\n",
+    "magnetism = 'NM'\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] * 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",
+    "        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 + f\"_{magnetism}\"))\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",
+    "        if magnetism == 'FM':\n",
+    "            if m == 0:\n",
+    "                continue\n",
+    "            magmoms = [m, m]\n",
+    "        elif magnetism == 'AFM':\n",
+    "            if m == 0:\n",
+    "                continue\n",
+    "            magmoms = [m, -m]\n",
+    "        elif magnetism == 'NM':\n",
+    "            magmoms = [0, 0]\n",
+    "            \n",
+    "        traj_fpath = out_dir / \"traj.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",
+    "        restart_fpath = out_dir / 'restart.gpw'\n",
+    "\n",
+    "        calc = GPAW(\n",
+    "            mode=PW(1000),\n",
+    "            xc='PBE',\n",
+    "            spinpol=True,\n",
+    "            # basis='dzp'\n",
+    "            basis='szp(dzp)',\n",
+    "            # h=0.25,\n",
+    "            # nbands=0 if element.is_noble_gas else '110%',\n",
+    "            hund=False,\n",
+    "            mixer=MixerDif(0.01, 1, 1) if element.is_transition_metal else MixerDif(0.25, 3, 10),\n",
+    "            eigensolver='cg', #'rmm-diis', #Davidson(3),  # This solver can parallelize over bands Davidson(3), #\n",
+    "            occupations=FermiDirac(0.0, fixmagmom=False), # if not element.is_metal else FermiDirac(0.2, fixmagmom=False),\n",
+    "            # eigensolver=LCAOETDM(),\n",
+    "            #     # searchdir_algo={'name': 'l-bfgs-p', 'memory': 10}),\n",
+    "            # occupations={'name': 'fixed-uniform'},\n",
+    "            # mixer={'backend': 'no-mixing'},\n",
+    "            # nbands='nao',\n",
+    "            symmetry={'point_group': False},\n",
+    "            txt=out_dir / 'out.txt',\n",
+    "            convergence={\n",
+    "                'eigenstates': 1e-5,\n",
+    "                'density': 5e-3,\n",
+    "                'energy': 5e-4,\n",
+    "                # 'bands': 4\n",
+    "            },\n",
+    "            # {'energy': 0.0005,  # eV / electron\n",
+    "            #  'density': 1.0e-4,  # electrons / electron\n",
+    "            #  'eigenstates': 4.0e-8,  # eV^2 / electron\n",
+    "            #  'bands': 'occupied'}\n",
+    "        )\n",
+    "        # calc.attach(calc.write, 10, restart_fpath, mode='all')\n",
+    "\n",
+    "        atoms.calc = calc\n",
+    "        \n",
+    "        # cdft = CDFT(calc=calc, atoms=atoms, spinspin_regions=      \n",
+    "        # atoms.calc = cdft\n",
+    "\n",
+    "        for i, r in enumerate(tqdm(np.flip(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",
+    "            # if i > 0:   \n",
+    "            #     magmoms = atoms.get_magnetic_moments()\n",
+    "                # m = min(abs(magmoms[0])*1.2, m)\n",
+    "                # magmoms = magmoms*m/np.abs(magmoms)\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",
+    "\n",
+    "\n",
+    "for symbol in tqdm(chemical_symbols):\n",
+    "    \n",
+    "    for magnetism in ['AFM', 'FM', 'NM']:\n",
+    "    \n",
+    "        da = symbol + symbol\n",
+    "\n",
+    "        # out_dir = Path(da)\n",
+    "        out_dir = Path(str(da + f\"_{magnetism}\"))\n",
+    "\n",
+    "        traj_fpath = out_dir / \"traj.extxyz\"\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': f'GGA-PBE (GPAW): {magnetism}',\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": "5d4a7312-a619-411f-9c6f-36c40cd47a34",
+   "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
+}

mlip_arena/tasks/diatomics/vasp/run.ipynb

@@ -0,0 +1,509 @@
+{
+ "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
+}