{
 "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
}