{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "154d45d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "eigvec = np.load(\"pbe/mode-1.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb7637b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from mlip_arena.models import MLIPEnum\n",
    "\n",
    "npoints = 101\n",
    "max_disp = 0.5\n",
    "disps = np.linspace(-max_disp, max_disp, npoints)\n",
    "\n",
    "replicas = (2, 2, 2)\n",
    "\n",
    "models = [\"MACE-MP(M)\", \"MatterSim\", \"ORBv2\", \"SevenNet\", \"CHGNet\", \"M3GNet\"]\n",
    "\n",
    "fig, axes = plt.subplots(\n",
    "    figsize=(6, 4),\n",
    "    nrows=2,\n",
    "    ncols=len(models) // 2,\n",
    "    constrained_layout=True,\n",
    ")\n",
    "\n",
    "axes = axes.flatten()\n",
    "\n",
    "i = 0\n",
    "\n",
    "for model in MLIPEnum:\n",
    "    if model.name not in models:\n",
    "        continue\n",
    "\n",
    "    calc = get_calculator(model, device=\"cuda\")\n",
    "\n",
    "    emin = float(\"inf\")\n",
    "    for a in [3.7, 3.85, 4.0, 4.15]:  # [3.8, 3.9, 4.0, 4.1]\n",
    "        atoms = read(\"BZO_cubic_prim.xyz\")\n",
    "        atoms.set_cell(cell=[a, a, a], scale_atoms=True)\n",
    "        # atoms = atoms * replicas\n",
    "\n",
    "        energies = []\n",
    "\n",
    "        for disp in disps:\n",
    "            atoms_disp = atoms.copy()\n",
    "            atoms_disp.calc = calc\n",
    "\n",
    "            atoms_disp.positions += eigvec * disp\n",
    "\n",
    "            energy = atoms_disp.get_potential_energy() / len(atoms_disp)\n",
    "            energies.append(energy)\n",
    "\n",
    "        energies = np.array(energies)\n",
    "\n",
    "        # shift the middle to 0\n",
    "        energies -= energies[npoints // 2]\n",
    "        emin = min(emin, energies.min())\n",
    "\n",
    "        axes[i].plot(disps, energies * 1000, label=f\"{a:.2f} Å\")\n",
    "\n",
    "    axes[i].axhline(0, color=\"black\", alpha=0.1)\n",
    "    emin -= 1e-4\n",
    "    axes[i].set(\n",
    "        title=model.name,\n",
    "        xlabel=\"Displacement (Å)\",\n",
    "        ylabel=\"Energy (meV/atom)\",\n",
    "        # xlim=(-0.1, 0.1),\n",
    "        # ylim=(-0.1, 0.5),\n",
    "    )\n",
    "\n",
    "    i += 1\n",
    "    # break\n",
    "\n",
    "fig.legend(\n",
    "    axes[0].get_legend_handles_labels()[0],\n",
    "    axes[0].get_legend_handles_labels()[1],\n",
    "    loc=\"lower center\",\n",
    "    bbox_to_anchor=(0.5, 1),\n",
    "    ncol=len(axes[0].get_legend_handles_labels()[0]),\n",
    "    # fontsize=6\n",
    ")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5eeb801",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}