Update notebook

Author: ae-foster

notebooks/02_bbmep/create_figures.ipynb

@@ -15,6 +15,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import sys\n",
+    "sys.path.append('/home/fosteradam/oneqmc/src')\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
     "from matplotlib import rcParams\n",
     "import numpy as np\n",
@@ -58,7 +61,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "save_figures = False"
+    "save_figures = True"
    ]
   },
   {
@@ -168,6 +171,7 @@
     "# Single color for all non-QMC methods, different markers\n",
     "non_qmc_color = \"#E05A33\"  # red/orange\n",
     "non_qmc_markers = ['s', '^', 'v', 'D', 'p', 'h', '*', 'X', 'P', '<', '>']  # variety of markers\n",
+    "non_qmc_marker_sizes = {'s': 180, '^': 240, 'v': 240, 'D': 180, 'p': 180, 'h': 180, '*': 300, 'X': 180, 'P': 180, '<': 180, '>': 180}\n",
     "\n",
     "fig, ax = plt.subplots(3, 2, figsize=(18.5, 19), sharey=True)\n",
     "compute_sec = {\n",
@@ -195,9 +199,9 @@
     "        x_gpu_hr = x * 4 * compute_sec[i] / 3600\n",
     "        y = 627.5 * MAE[i][j][:, 1]\n",
     "        h, v = np.divmod(m, 3)\n",
-    "        ax[v][h].scatter(x_gpu_hr, y, marker=\"o\", s=140, label=j, color=qmc_colors[n])\n",
+    "        ax[v][h].scatter(x_gpu_hr, y, marker=\"o\", s=140, label=j, color=qmc_colors[n], zorder=3)\n",
     "        a, b = np.polyfit(np.log(x_gpu_hr), np.log(y), 1)\n",
-    "        ax[v][h].plot(x_gpu_hr, np.exp(a * np.log(x_gpu_hr) + b), color=qmc_colors[n], linewidth=4)\n",
+    "        ax[v][h].plot(x_gpu_hr, np.exp(a * np.log(x_gpu_hr) + b), color=qmc_colors[n], linewidth=4, zorder=2)\n",
     "        ax[v][h].set_xscale(\"log\", base=10)\n",
     "        ax[v][h].set_yscale(\"log\", base=4)\n",
     "    ticks = [0.01, 0.1, 1, 10, 100]\n",
@@ -210,21 +214,25 @@
     "            gpu_hr_equiv,\n",
     "            627.5 * MAE_es[i][k][1],\n",
     "            label=k,\n",
-    "            s=180,\n",
+    "            s=non_qmc_marker_sizes[non_qmc_markers[idx % len(non_qmc_markers)]],\n",
     "            color=non_qmc_color,\n",
     "            marker=non_qmc_markers[idx % len(non_qmc_markers)],\n",
+    "            zorder=4,\n",
+    "            edgecolors='white',\n",
+    "            linewidths=1.5,\n",
     "        )\n",
     "    ax[v][h].set_xticks(ticks)\n",
     "    ax[v][h].set_xticklabels(ticks, fontsize=16)\n",
-    "    ax[v][h].axhspan(0.2, 5, color=\"grey\", alpha=0.2, lw=0)\n",
-    "    ax[v][h].set_xlim([0.0025, 500])\n",
-    "    ax[v][h].axhline(1, color=\"k\", ls=\":\", linewidth=4)\n",
+    "    ax[v][h].axhspan(0.2, 5, color=\"grey\", alpha=0.2, lw=0, zorder=0)\n",
+    "    ax[v][h].set_xlim([0.0025, 300] if h == 0 else [0.005, 1100])\n",
+    "    ax[v][h].axhline(1, color=\"darkgray\", ls=\":\", linewidth=4, zorder=1)\n",
     "    ax[v][h].set_yticks(\n",
-    "        [0.5, 1, 2, 5, 10, 20, 50], [0.5, 1, 2, 5, 10, 20, 50], fontsize=16\n",
+    "        [0.5, 1, 2, 5, 10, 20, 50, 100, 200], [0.5, 1, 2, 5, 10, 20, 50, 100, 200], fontsize=16\n",
     "    )\n",
+    "    ax[v][h].minorticks_off()\n",
     "    ax[v][h].set_ylim([0.25, 850])\n",
     "    if h == 0:\n",
-    "        ax[v][h].set_ylabel(\"MARE over entire MEP (kcal/mol)\", fontsize=22)\n",
+    "        ax[v][h].set_ylabel(\"MARE over MEP (kcal/mol)\", fontsize=22)\n",
     "ax[2][1].set_axis_off()\n",
     "\n",
     "ax[0][0].legend(\n",
@@ -335,7 +343,7 @@
    "source": [
     "## Get reaction profile figure\n",
     "\n",
-    "fig, ax = plt.subplots(3, 2, figsize=(13, 17.5), sharex=True)\n",
+    "fig, ax = plt.subplots(3, 2, figsize=(13, 19), sharex=True)\n",
     "\n",
     "for m, j in enumerate([\"Ethane\", \"Formamide\", \"2-Aminopropan-2-ol\"]):\n",
     "    step = int(MAE[j][\"Orbformer fine-tune LAC (400k), all structures\"][-1, 0])\n",
@@ -440,8 +448,8 @@
     "        )\n",
     "    ax[m][0].set_ylabel(\"Relative E (kcal/mol)\", fontsize=22)\n",
     "    ax[m][1].set_ylabel(\"Relative E Error (kcal/mol)\", fontsize=22)\n",
-    "ax[0][0].text(0.03, 0.9, \"(a) Ethane\", transform=ax[0][0].transAxes, fontsize=24)\n",
-    "ax[0][1].text(0.03, 0.9, \"(d) Ethane\", transform=ax[0][1].transAxes, fontsize=24)\n",
+    "ax[0][0].text(0.03, 1.02, \"(a) Ethane\", transform=ax[0][0].transAxes, fontsize=24)\n",
+    "ax[0][1].text(0.03, 1.02, \"(d) Ethane\", transform=ax[0][1].transAxes, fontsize=24)\n",
     "ax[1][0].text(0.03, 0.9, \"(b) Formamide\", transform=ax[1][0].transAxes, fontsize=24)\n",
     "ax[1][1].text(0.03, 0.9, \"(e) Formamide\", transform=ax[1][1].transAxes, fontsize=24)\n",
     "ax[2][0].text(\n",
@@ -508,14 +516,18 @@
     "ax[0][1].text(\n",
     "    0.18, 1.1, \"Relative Energy Error\", transform=ax[0][1].transAxes, fontsize=26\n",
     ")\n",
+    "ax[0][0].set_ylim([-100, 55])\n",
     "ax[2][0].set_ylim([-65, 58])\n",
-    "plt.subplots_adjust(wspace=0.2, hspace=0.02)\n",
+    "fig.align_ylabels(ax[:, 0])\n",
+    "fig.align_ylabels(ax[:, 1])\n",
+    "plt.subplots_adjust(wspace=0.2, hspace=0.15)\n",
     "ax[2][0].set_xticks(2 * np.arange(11))\n",
     "ax[2][0].set_xlim([-0.5, 19.5])\n",
     "ax[2][0].set_xticklabels(2 * np.arange(11))\n",
     "ax[2][0].set_xlabel(\"Image ID on MEP\", fontsize=22)\n",
     "ax[2][1].set_xlabel(\"Image ID on MEP\", fontsize=22)\n",
     "plt.draw()\n",
+    "\n",
     "if save_figures:\n",
     "    plt.savefig(f\"{data_dir}/BBMEP_energy_profile.pdf\", dpi=600, bbox_inches='tight')\n",
     "plt.show()"
@@ -532,7 +544,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "oneqmc",
    "language": "python",
    "name": "python3"
   },