add alignn

.gitignore

@@ -1,10 +1,11 @@
+tests/
+*.out
+mlip_arena/tasks/*/*/*/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
 *$py.class
-tests/
-*.out
-mlip_arena/tasks/*/*/
 
 # C extensions
 *.so

mlip_arena/tasks/diatomics/m3gnet/run.ipynb

@@ -0,0 +1,295 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "3200850a-b8fb-4f50-9815-16ae8da0f942",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "Do you really want to delete everything in /global/homes/c/cyrusyc/.cache/matgl (y|n)?  y\n"
+     ]
+    },
+    {
+     "ename": "ValueError",
+     "evalue": "Bad serialized model or bad model name. It is possible that you have an older model cached. Please clear your cache by running `python -c \"import matgl; matgl.clear_cache()\"`",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/matgl/utils/io.py:212\u001b[0m, in \u001b[0;36mload_model\u001b[0;34m(path, **kwargs)\u001b[0m\n\u001b[1;32m    211\u001b[0m         cls_ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(mod, classname)\n\u001b[0;32m--> 212\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcls_\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfpaths\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    213\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/matgl/utils/io.py:129\u001b[0m, in \u001b[0;36mIOMixIn.load\u001b[0;34m(cls, path, **kwargs)\u001b[0m\n\u001b[1;32m    128\u001b[0m d \u001b[38;5;241m=\u001b[39m {k: v \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m d\u001b[38;5;241m.\u001b[39mitems() \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m k\u001b[38;5;241m.\u001b[39mstartswith(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m@\u001b[39m\u001b[38;5;124m\"\u001b[39m)}\n\u001b[0;32m--> 129\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43md\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m model\u001b[38;5;241m.\u001b[39mload_state_dict(state, strict\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n",
+      "\u001b[0;31mTypeError\u001b[0m: Potential.__init__() got an unexpected keyword argument 'calc_magmom'",
+      "\nThe above exception was the direct cause of the following exception:\n",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[6], line 18\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[1;32m     17\u001b[0m matgl\u001b[38;5;241m.\u001b[39mclear_cache()\n\u001b[0;32m---> 18\u001b[0m potential \u001b[38;5;241m=\u001b[39m \u001b[43mmatgl\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_model\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mM3GNet-MP-2021.2.8-PES\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     19\u001b[0m calculator \u001b[38;5;241m=\u001b[39m PESCalculator(potential)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/matgl/utils/io.py:214\u001b[0m, in \u001b[0;36mload_model\u001b[0;34m(path, **kwargs)\u001b[0m\n\u001b[1;32m    212\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m cls_\u001b[38;5;241m.\u001b[39mload(fpaths, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    213\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[0;32m--> 214\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    215\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBad serialized model or bad model name. It is possible that you have an older model cached. Please \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    216\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mclear your cache by running `python -c \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mimport matgl; matgl.clear_cache()\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    217\u001b[0m     ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n",
+      "\u001b[0;31mValueError\u001b[0m: Bad serialized model or bad model name. It is possible that you have an older model cached. Please clear your cache by running `python -c \"import matgl; matgl.clear_cache()\"`"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import numpy as np\n",
+    "\n",
+    "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",
+    "from pymatgen.core import Element\n",
+    "import pandas as pd\n",
+    "\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "import matgl\n",
+    "from matgl.ext.ase import PESCalculator\n",
+    "\n",
+    "matgl.clear_cache()\n",
+    "potential = matgl.load_model(\"M3GNet-MP-2021.2.8-PES\")\n",
+    "calculator = PESCalculator(potential)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "90887faa-1601-4c4c-9c44-d16731471d7f",
+   "metadata": {
+    "scrolled": true,
+    "tags": []
+   },
+   "outputs": [
+    {
+     "ename": "ValueError",
+     "evalue": "Bad serialized model or bad model name. It is possible that you have an older model cached. Please clear your cache by running `python -c \"import matgl; matgl.clear_cache()\"`",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/matgl/utils/io.py:212\u001b[0m, in \u001b[0;36mload_model\u001b[0;34m(path, **kwargs)\u001b[0m\n\u001b[1;32m    211\u001b[0m         cls_ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(mod, classname)\n\u001b[0;32m--> 212\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcls_\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfpaths\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    213\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/matgl/utils/io.py:129\u001b[0m, in \u001b[0;36mIOMixIn.load\u001b[0;34m(cls, path, **kwargs)\u001b[0m\n\u001b[1;32m    128\u001b[0m d \u001b[38;5;241m=\u001b[39m {k: v \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m d\u001b[38;5;241m.\u001b[39mitems() \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m k\u001b[38;5;241m.\u001b[39mstartswith(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m@\u001b[39m\u001b[38;5;124m\"\u001b[39m)}\n\u001b[0;32m--> 129\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43md\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m model\u001b[38;5;241m.\u001b[39mload_state_dict(state, strict\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n",
+      "\u001b[0;31mTypeError\u001b[0m: Potential.__init__() got an unexpected keyword argument 'calc_magmom'",
+      "\nThe above exception was the direct cause of the following exception:\n",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[2], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m potential \u001b[38;5;241m=\u001b[39m \u001b[43mmatgl\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_model\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mM3GNet-MP-2021.2.8-PES\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      2\u001b[0m calculator \u001b[38;5;241m=\u001b[39m PESCalculator(potential)\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m symbol \u001b[38;5;129;01min\u001b[39;00m tqdm(chemical_symbols):\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/matgl/utils/io.py:214\u001b[0m, in \u001b[0;36mload_model\u001b[0;34m(path, **kwargs)\u001b[0m\n\u001b[1;32m    212\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m cls_\u001b[38;5;241m.\u001b[39mload(fpaths, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    213\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[0;32m--> 214\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    215\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBad serialized model or bad model name. It is possible that you have an older model cached. Please \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    216\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mclear your cache by running `python -c \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mimport matgl; matgl.clear_cache()\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    217\u001b[0m     ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n",
+      "\u001b[0;31mValueError\u001b[0m: Bad serialized model or bad model name. It is possible that you have an older model cached. Please clear your cache by running `python -c \"import matgl; matgl.clear_cache()\"`"
+     ]
+    }
+   ],
+   "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 / \"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",
+    "        calc = calculator\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",
+    "            # 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": 2,
+   "id": "a0ac2c09-370b-4fdd-bf74-ea5c4ade0215",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cc766db4ce844c40848791e14a71832c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/119 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "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 / \"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': 'M3GNet',\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": "2207f50e-63a1-4199-b2e1-a11858af5108",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df"
+   ]
+  }
+ ],
+ "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
+}