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