diff --git a/notebooks/02_bbmep/create_figures.ipynb b/notebooks/02_bbmep/create_figures.ipynb index 2464ca3..caaa199 100644 --- a/notebooks/02_bbmep/create_figures.ipynb +++ b/notebooks/02_bbmep/create_figures.ipynb @@ -22,11 +22,16 @@ "from matplotlib.offsetbox import OffsetImage, AnnotationBbox\n", "import matplotlib.image as mpimg\n", "from scipy.ndimage import rotate\n", + "from matplotlib.ticker import NullLocator\n", + "\n", "\n", "from oneqmc.analysis.plot import set_defaults\n", "from oneqmc.analysis import colours\n", "\n", "set_defaults()\n", + "rcParams[\"axes.spines.top\"] = False\n", + "rcParams['xtick.labelsize'] = 20\n", + "rcParams['ytick.labelsize'] = 20\n", "\n", "colormaps_es = {\n", " \"UM06-2X-D3/Dev2-QZVPP\": colours.BRAND_BLUE,\n", @@ -88,10 +93,17 @@ "\n", "def remove_black_edge(img: np.ndarray):\n", " # background is set to 0.0 opacity\n", - " background_idxs = np.where(img[:,:,3] < 1e-6)\n", + " background_idxs = np.where(img[:, :, 3] < 1e-6)\n", " # set color of background to white\n", " img[background_idxs[0], background_idxs[1], :3] = 1.0\n", - " return img" + " return img\n", + "\n", + "\n", + "def remove_background(img: np.ndarray, white_thresh: float = 0.999):\n", + " # sets ~white pixels of an image to 0.0 opacity\n", + " mask = np.all(img[:, :, :3] > white_thresh, axis=2)\n", + " img[mask, 3] = 0\n", + " return np.clip(img, 0, 1)" ] }, { @@ -139,18 +151,27 @@ " ]\n", " MAE[i] = dict.fromkeys(method, {})\n", " for j in method:\n", - " steps = sorted(\n", - " set(orb_energies.loc[orb_energies[\"Ansatz\"] == j][\"Fine-tune/Train_steps\"])\n", - " )\n", + " steps = sorted(set(orb_energies.loc[orb_energies[\"Ansatz\"] == j][\"Fine-tune/Train_steps\"]))\n", " MAE[i][j] = []\n", " for k in steps:\n", " a = orb_energies.loc[orb_energies[\"Ansatz\"] == j]\n", - " e = np.array(\n", - " (a.loc[orb_energies[\"Fine-tune/Train_steps\"] == k][\"E(Hartree)\"])\n", - " )\n", + " e = np.array((a.loc[orb_energies[\"Fine-tune/Train_steps\"] == k][\"E(Hartree)\"]))\n", " mae = np.mean(np.abs(relative(e) - relative(reference_e)))\n", " MAE[i][j].append([k, mae])\n", - " MAE[i][j] = np.vstack(MAE[i][j])" + " MAE[i][j] = np.vstack(MAE[i][j])\n", + "\n", + "compute_sec = {\n", + " \"Ethane\": 16265.031 / 32000,\n", + " \"Formamide\": 10323.08825 / 32000,\n", + " \"1-Propanol\": 85435.938 / 80000,\n", + " \"2-Aminopropan-2-ol\": 92369.6905 / 80000,\n", + " \"L-Alanine\": 160420.138 / 100000,\n", + "} # 4 GPU training costs\n", + "# Psiformer 2x timing 296013.404/200000 for single GPU, Psiformer 4x timing 178873.479/200000 for 4 GPU\n", + "psiformer_factor = {\n", + " \"Ethane\": 296013.404 / 200000 / (4 * 16265.031 / 32000),\n", + " \"Formamide\": 182497.587 / 200000 / (10323.08825 / 32000),\n", + "}" ] }, { @@ -160,194 +181,216 @@ "metadata": {}, "outputs": [], "source": [ - "cost_option = \"hour\" #'wallclock_hour', 'energy', 'hour'\n", "# {\"Orbformer scratch, single point\": '#74baae', \"Orbformer scratch, all points\": '#4e938d',\n", "# \"Orbformer finetune LAC (200k), all points\": '#366c6c', \"Orbformer finetune LAC (400k), all points\":'#27474a',\n", "# \"Orbformer finetune LAC (1000k), all points\":'#1d2429', \"Psiformer scratch, single point\":'#8de971'}\n", - "qmc_colors = [\"#74baae\", \"#4e938d\", \"#27474a\", \"#8de971\"]\n", "\n", "\n", - "fig, ax = plt.subplots(3, 2, figsize=(18.5, 19), sharey=True)\n", - "compute_sec = {\n", - " \"Ethane\": 16265.031 / 32000,\n", - " \"Formamide\": 10323.08825 / 32000,\n", - " \"1-Propanol\": 85435.938 / 80000,\n", - " \"2-Aminopropan-2-ol\": 92369.6905 / 80000,\n", - " \"L-Alanine\": 160420.138 / 100000,\n", - "} # 4 GPU training costs\n", - "# Psiformer 2x timing 296013.404/200000 for single GPU, Psiformer 4x timing 178873.479/200000 for 4 GPU\n", - "psiformer_factor = {\n", - " \"Ethane\": 296013.404 / 200000 / (4 * 16265.031 / 32000),\n", - " \"Formamide\": 182497.587 / 200000 / (10323.08825 / 32000),\n", - "}\n", - "rcParams[\"axes.spines.top\"] = True\n", - "ax_top = {}\n", - "print(MAE.keys())\n", - "for m, i in enumerate(MAE.keys()):\n", - " for n, j in enumerate(MAE[i].keys()):\n", - " x = MAE[i][j][:, 0]\n", - " if \"Psiformer\" in j.split():\n", - " x = x * psiformer_factor[i]\n", - " if \"all\" in j.split():\n", - " x = x / 20\n", - " y = 627.5 * MAE[i][j][:, 1]\n", - " h, v = np.divmod(m, 3)\n", - " ax[v][h].scatter(x, y, marker=\"o\", s=140, label=j, color=qmc_colors[n])\n", - " a, b = np.polyfit(np.log(x), np.log(y), 1)\n", - " ax[v][h].plot(x, np.exp(a * np.log(x) + b), color=qmc_colors[n], linewidth=4)\n", - " ax[v][h].set_xscale(\"log\", base=4)\n", - " ax[v][h].set_yscale(\"log\", base=4)\n", - " ticks = [10, 50, 250, 1000, 4000, 16000, 64000, 300000]\n", - " cost_factor = 300 / (280 / 64)\n", - " for k in MAE_es[i].keys():\n", - " ax[v][h].scatter(\n", - " MAE_es[i][k][0] / (4 * compute_sec[i] / 3600 * cost_factor),\n", - " 627.5 * MAE_es[i][k][1],\n", - " label=k,\n", - " s=180,\n", - " color=colormaps_es[k],\n", - " marker=marker_es[k],\n", - " )\n", - " ax[v][h].set_xticks(ticks, ticks, fontsize=16)\n", - " ax[v][h].axhspan(0.2, 5, color=\"grey\", alpha=0.2, lw=0)\n", - " ax[v][h].set_xlim([6, 580000])\n", - " ax_top[m] = ax[v][h].twiny()\n", - " ax_top[m].set_xscale(\"log\", base=4)\n", - " ax_top[m].set_xlim([6, 580000])\n", - " if cost_option == \"hour\":\n", - " ax[v][h].set_xticklabels(\n", - " np.round(np.array(ticks) * 4 * compute_sec[i] / 3600 * cost_factor, 1)\n", - " )\n", - " else:\n", - " ax[v][h].set_xticklabels(cost)\n", - " ax[v][h].axhline(1, color=\"k\", ls=\":\", linewidth=4)\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", + "def make_annotation_box(arrimg, xy, zoom, mols_offset):\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=zoom)\n", + " ab = AnnotationBbox(\n", + " imagebox,\n", + " tuple(x + y for x, y in zip(xy, mols_offset)),\n", + " xycoords=\"axes fraction\",\n", + " frameon=False,\n", + " clip_on=False,\n", " )\n", - " ax[v][h].set_ylim([0.2, 99])\n", - " if v == 0:\n", - " if cost_option == \"hour\":\n", - " ax_top[m].set_xlabel(\"Orbformer A100 GPU hr/structure\", fontsize=22)\n", - " else:\n", - " ax_top[m].set_xlabel(\n", - " \"Scratch training or finetuning steps per structure\", fontsize=22\n", + " return ab\n", + "\n", + "\n", + "def make_scaling(fig, ax, legend_offset=(0, 0), mols_offset=(0, 0)):\n", + " qmc_colors = [\"#74baae\", \"#4e938d\", \"#27474a\", \"#8de971\"]\n", + " # 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\", \"<\", \">\"]\n", + " non_qmc_marker_sizes = [180, 180, 180, 150, 180, 180, 400, 180, 180, 180, 180]\n", + "\n", + " for m, i in enumerate(MAE.keys()):\n", + " for n, j in enumerate(MAE[i].keys()):\n", + " x = MAE[i][j][:, 0]\n", + " if \"Psiformer\" in j.split():\n", + " x = x * psiformer_factor[i]\n", + " if \"all\" in j.split():\n", + " x = x / 20\n", + " # Convert to GPU hours (x is already in GPU-relative units)\n", + " 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", + " clip = y <= 250\n", + " ax[v][h].scatter(\n", + " x_gpu_hr[clip],\n", + " y[clip],\n", + " marker=\"o\",\n", + " s=140,\n", + " label=j,\n", + " color=qmc_colors[n],\n", + " zorder=50,\n", + " edgecolors=(1, 1, 1, 0.5),\n", + " linewidths=2,\n", + " )\n", + " a, b = np.polyfit(np.log(x_gpu_hr), np.log(y), 1)\n", + " ax[v][h].plot(\n", + " x_gpu_hr[clip],\n", + " np.exp(a * np.log(x_gpu_hr[clip]) + b),\n", + " color=qmc_colors[n],\n", + " linewidth=4,\n", + " zorder=40,\n", + " )\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, 1000]\n", + " cost_factor = 300 / (280 / 64)\n", + " for idx, k in enumerate(MAE_es[i].keys()):\n", + " # Convert CPU cost to equivalent GPU hours for plotting\n", + " cpu_hr = MAE_es[i][k][0]\n", + " gpu_hr_equiv = cpu_hr / cost_factor # convert CPU hr to GPU hr equivalent\n", + " ax[v][h].scatter(\n", + " gpu_hr_equiv,\n", + " 627.5 * MAE_es[i][k][1],\n", + " label=k,\n", + " s=non_qmc_marker_sizes[idx % len(non_qmc_marker_sizes)],\n", + " color=non_qmc_color,\n", + " marker=non_qmc_markers[idx % len(non_qmc_markers)],\n", + " zorder=100,\n", + " edgecolors=(1, 1, 1, 0.5),\n", + " linewidths=2,\n", " )\n", - " if h == 0:\n", - " ax[v][h].set_ylabel(\"MARE over entire MEP (kcal/mol)\", fontsize=22)\n", - " if cost_option == \"hour\":\n", - " ax_top[m].set_xticks(ticks)\n", - " ax_top[m].set_xticklabels(\n", - " np.round(np.array(ticks) * 4 * compute_sec[i] / 3600, 2), fontsize=16\n", + " ax[v][h].set_xticks(ticks)\n", + " if (v, h) == (2, 0) or (v, h) == (1, 1):\n", + " ax[v][h].set_xticklabels(ticks)\n", + " else:\n", + " ax[v][h].set_xticklabels([])\n", + " ax[v][h].axhspan(0.2, 5, color=\"grey\", alpha=0.2, lw=0)\n", + " ax[v][h].set_xlim([0.0025, 1400])\n", + " ax[v][h].axhline(1, color=\"k\", ls=\":\", linewidth=3, zorder=0)\n", + " ax[v][h].set_yticks(\n", + " [0.5, 1, 2, 5, 10, 20, 50, 100, 200], [0.5, 1, 2, 5, 10, 20, 50, 100, 200],\n", " )\n", - " else:\n", - " ax_top[m].set_xticks(ticks)\n", - " ax_top[m].set_xticklabels(ticks)\n", - "ax[2][1].set_axis_off()\n", "\n", - "ax[0][0].legend(\n", - " bbox_to_anchor=(2.06, -1.28),\n", - " fontsize=19,\n", - " handletextpad=0.02,\n", - " borderpad=0.16,\n", - " labelspacing=0.4,\n", - ")\n", - "ax[0][0].text(0.03, 0.9, \"(g) Ethane\", transform=ax[0][0].transAxes, fontsize=24)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_0.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.12)\n", - "ab = AnnotationBbox(imagebox, (2500, 20), frameon=False)\n", - "ax[0][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_19.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.2)\n", - "ab = AnnotationBbox(imagebox, (35000, 15), frameon=False)\n", - "ax[0][0].add_artist(ab)\n", - "ax[0][0].annotate(\n", - " \"\", xytext=(5500, 20), xy=(12000, 20), arrowprops=dict(arrowstyle=\"->\")\n", - ")\n", + " ax[v][h].yaxis.set_minor_locator(NullLocator())\n", + " ax[v][h].set_ylim([0.25, 1000])\n", + " if h == 0:\n", + " ax[v][h].set_ylabel(\"MARE over MEP (kcal/mol)\", fontsize=22)\n", + " ax[2][1].set_axis_off()\n", "\n", - "ax[1][0].text(0.03, 0.9, \"(h) Formamide\", transform=ax[1][0].transAxes, fontsize=24)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_0.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.12)\n", - "ab = AnnotationBbox(imagebox, (10000, 30), frameon=False)\n", - "ax[1][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_19.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.2)\n", - "ab = AnnotationBbox(imagebox, (100000, 25), frameon=False)\n", - "ax[1][0].add_artist(ab)\n", - "ax[1][0].annotate(\n", - " \"\", xytext=(22000, 30), xy=(45000, 30), arrowprops=dict(arrowstyle=\"->\")\n", - ")\n", + " ax[0][0].legend(\n", + " bbox_to_anchor=tuple(x + y for x, y in zip((2.06, -1.28), legend_offset)),\n", + " fontsize=19,\n", + " handletextpad=0.02,\n", + " borderpad=0.16,\n", + " labelspacing=0.4,\n", + " )\n", + " ax[0][0].text(0.03, 0.9, \"(g) Ethane\", transform=ax[0][0].transAxes, fontsize=24, zorder=200)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_0.png\")\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.62, 0.88), zoom=0.12, mols_offset=mols_offset)\n", + " ax[0][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_19.png\")\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.88, 0.82), zoom=0.2, mols_offset=mols_offset)\n", + " ax[0][0].add_artist(ab)\n", + " ax[0][0].annotate(\n", + " \"\",\n", + " xytext=tuple(x + y for x, y in zip((0.71, 0.88), mols_offset)),\n", + " xy=tuple(x + y for x, y in zip((0.77, 0.88), mols_offset)),\n", + " xycoords=\"axes fraction\",\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", "\n", - "ax[2][0].text(0.03, 0.9, \"(i) 1-Propanol\", transform=ax[2][0].transAxes, fontsize=24)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/1-Propanol/1-Propanol_0.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.17)\n", - "ab = AnnotationBbox(imagebox, (12000, 38), frameon=False)\n", - "ax[2][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/1-Propanol/1-Propanol_19.png\")\n", - "arrimg = np.rot90(arrimg)\n", - "imagebox = OffsetImage(arrimg, zoom=0.25)\n", - "ab = AnnotationBbox(imagebox, (120000, 28), frameon=False)\n", - "ax[2][0].add_artist(ab)\n", - "ax[2][0].annotate(\n", - " \"\", xytext=(30000, 38), xy=(55000, 38), arrowprops=dict(arrowstyle=\"->\")\n", - ")\n", + " ax[1][0].text(0.03, 0.9, \"(h) Formamide\", transform=ax[1][0].transAxes, fontsize=24, zorder=200)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_0.png\")\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.68, 0.88), zoom=0.1, mols_offset=mols_offset)\n", + " ax[1][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_19.png\")\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.92, 0.82), zoom=0.18, mols_offset=mols_offset)\n", + " ax[1][0].add_artist(ab)\n", + " ax[1][0].annotate(\n", + " \"\",\n", + " xytext=tuple(x + y for x, y in zip((0.77, 0.88), mols_offset)),\n", + " xy=tuple(x + y for x, y in zip((0.83, 0.88), mols_offset)),\n", + " xycoords=\"axes fraction\",\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", "\n", - "ax[0][1].text(\n", - " 0.03, 0.9, \"(j) 2-Aminopropan-2-ol\", transform=ax[0][1].transAxes, fontsize=24\n", - ")\n", - "arrimg = mpimg.imread(\n", - " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_0.png\"\n", - ")\n", - "imagebox = OffsetImage(arrimg, zoom=0.18)\n", - "ab = AnnotationBbox(imagebox, (14000, 35), frameon=False)\n", - "ax[0][1].add_artist(ab)\n", - "arrimg = mpimg.imread(\n", - " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_19.png\"\n", - ")\n", - "arrimg = remove_black_edge(rotate(arrimg, angle=60, reshape=True))\n", - "imagebox = OffsetImage(arrimg, zoom=0.23)\n", - "ab = AnnotationBbox(imagebox, (120000, 28), frameon=False)\n", - "ax[0][1].add_artist(ab)\n", - "ax[0][1].annotate(\n", - " \"\", xytext=(28000, 30), xy=(60000, 30), arrowprops=dict(arrowstyle=\"->\")\n", - ")\n", + " ax[2][0].text(\n", + " 0.03, 0.9, \"(i) 1-Propanol\", transform=ax[2][0].transAxes, fontsize=24, zorder=200\n", + " )\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/1-Propanol/1-Propanol_0.png\")\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.68, 0.88), zoom=0.17, mols_offset=mols_offset)\n", + " ax[2][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/1-Propanol/1-Propanol_19.png\")\n", + " arrimg = np.rot90(arrimg)\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.92, 0.82), zoom=0.25, mols_offset=mols_offset)\n", + " ax[2][0].add_artist(ab)\n", + " ax[2][0].annotate(\n", + " \"\",\n", + " xytext=tuple(x + y for x, y in zip((0.77, 0.88), mols_offset)),\n", + " xy=tuple(x + y for x, y in zip((0.83, 0.88), mols_offset)),\n", + " xycoords=\"axes fraction\",\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", "\n", - "ax[1][1].text(0.03, 0.9, \"(k) L-Alanine\", transform=ax[1][1].transAxes, fontsize=24)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/L-Alanine/L-Alanine_0.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.18)\n", - "ab = AnnotationBbox(imagebox, (28000, 40), frameon=False)\n", - "ax[1][1].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/L-Alanine/L-Alanine_19.png\")\n", - "arrimg = remove_black_edge(rotate(arrimg, angle=60))\n", - "imagebox = OffsetImage(arrimg, zoom=0.16)\n", - "ab = AnnotationBbox(imagebox, (140000, 25), frameon=False)\n", - "ax[1][1].add_artist(ab)\n", - "ax[1][1].annotate(\n", - " \"\", xytext=(44000, 30), xy=(85000, 30), arrowprops=dict(arrowstyle=\"->\")\n", - ")\n", + " ax[0][1].text(\n", + " 0.03, 0.9, \"(j) 2-Aminopropan-2-ol\", transform=ax[0][1].transAxes, fontsize=24, zorder=200\n", + " )\n", + " arrimg = mpimg.imread(\n", + " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_0.png\"\n", + " )\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.68, 0.88), zoom=0.18, mols_offset=mols_offset)\n", + " ax[0][1].add_artist(ab)\n", + " arrimg = mpimg.imread(\n", + " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_19.png\"\n", + " )\n", + " arrimg = remove_black_edge(rotate(arrimg, angle=60, reshape=True))\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.92, 0.82), zoom=0.23, mols_offset=mols_offset)\n", + " ax[0][1].add_artist(ab)\n", + " ax[0][1].annotate(\n", + " \"\",\n", + " xytext=tuple(x + y for x, y in zip((0.77, 0.88), mols_offset)),\n", + " xy=tuple(x + y for x, y in zip((0.83, 0.88), mols_offset)),\n", + " xycoords=\"axes fraction\",\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", "\n", - "if cost_option == \"wallclock_hour\":\n", - " ax[2][0].set_xlabel(\"GPU or CPU hr/structure\", fontsize=22)\n", - " ax[1][1].set_xlabel(\"GPU or CPU hr/structure\", fontsize=22)\n", - "elif cost_option == \"energy\":\n", - " ax[2][0].set_xlabel(\"Energy cost/structure (KWh)\", fontsize=22)\n", - " ax[1][1].set_xlabel(\"Energy cost/structure (KWh)\", fontsize=22)\n", - "else:\n", - " ax[2][0].set_xlabel(\"Other methods AMD EPYC 7763 CPU hr/structure\", fontsize=22)\n", - " ax[1][1].set_xlabel(\"Other methods AMD EPYC 7763 CPU hr/structure\", fontsize=22)\n", + " ax[1][1].text(0.03, 0.9, \"(k) L-Alanine\", transform=ax[1][1].transAxes, fontsize=24, zorder=200)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/L-Alanine/L-Alanine_0.png\")\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.7, 0.88), zoom=0.18, mols_offset=mols_offset)\n", + " ax[1][1].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/L-Alanine/L-Alanine_19.png\")\n", + " arrimg = remove_black_edge(rotate(arrimg, angle=60))\n", + " ab = make_annotation_box(arrimg=arrimg, xy=(0.92, 0.82), zoom=0.16, mols_offset=mols_offset)\n", + " ax[1][1].add_artist(ab)\n", + " ax[1][1].annotate(\n", + " \"\",\n", + " xytext=tuple(x + y for x, y in zip((0.79, 0.88), mols_offset)),\n", + " xy=tuple(x + y for x, y in zip((0.85, 0.88), mols_offset)),\n", + " xycoords=\"axes fraction\",\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", "\n", - "ax[0][0].text(\n", - " 0.55,\n", - " 1.16,\n", - " \"Mean Absolute Relative Energy Error (MARE)\",\n", - " transform=ax[0][0].transAxes,\n", - " fontsize=26,\n", - ")\n", + " ax[2][0].set_xlabel(\"A100 GPU hr/structure\", fontsize=22)\n", + " ax[1][1].set_xlabel(\"A100 GPU hr/structure\", fontsize=22)\n", + "\n", + " ax[0][0].text(\n", + " 0.55,\n", + " 1.03,\n", + " \"Mean Absolute Relative Energy Error (MARE)\",\n", + " transform=ax[0][0].transAxes,\n", + " fontsize=26,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e0ab2c26", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(3, 2, figsize=(18.5, 19), sharey=True)\n", "\n", + "make_scaling(fig, ax)\n", "plt.subplots_adjust(wspace=0.04, hspace=0.15)\n", - "# plt.draw()\n", "if save_figures:\n", - " plt.savefig(f\"{data_dir}/BBMEP_result_fig_hours.pdf\", dpi=600, bbox_inches='tight')\n", - "# plt.show()" + " plt.savefig(f\"{data_dir}/BBMEP_result_fig_hours.pdf\", dpi=600, bbox_inches=\"tight\")" ] }, { @@ -359,189 +402,208 @@ "source": [ "## Get reaction profile figure\n", "\n", - "fig, ax = plt.subplots(3, 2, figsize=(13, 17.5), 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", - " qmc_e = raw_energy.loc[\n", - " (raw_energy[\"MEP_system\"] == j)\n", - " & (raw_energy[\"Ansatz\"] == \"Orbformer fine-tune LAC (400k), all structures\")\n", - " & (raw_energy[\"Fine-tune/Train_steps\"] == step)\n", - " ][\"E(Hartree)\"]\n", - " qmc_e = relative(np.array(qmc_e))\n", - " es_e = es_energies.loc[\n", - " (es_energies[\"MEP_system\"] == j)\n", - " & (es_energies[\"Method\"] == \"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\")\n", - " ][\"E(Hartree)\"]\n", - " es_e = relative(np.array(es_e))\n", - " ref_e = relative(np.array(ref.loc[ref[\"MEP_system\"] == j][\"Reference_E(Hartree)\"]))\n", - " std = np.array(ref.loc[ref[\"MEP_system\"] == j][\"Reference_E_std(Hartree)\"])\n", - " ax[m][0].plot(\n", - " range(20),\n", - " 627.5 * qmc_e,\n", - " color=\"#27474a\",\n", - " marker=\"o\",\n", - " label=\"Best Orbformer fine-tune LAC (400k), all structures\",\n", - " linewidth=3.5,\n", - " markersize=14,\n", - " linestyle=\"-.\",\n", - " )\n", - " ax[m][1].plot(\n", - " range(20),\n", - " 627.5 * (qmc_e - ref_e),\n", - " color=\"#27474a\",\n", - " marker=\"o\",\n", - " label=\"Best Orbformer fine-tune LAC (400k), all structures\",\n", - " linewidth=3.5,\n", - " markersize=14,\n", - " linestyle=\"-.\",\n", - " )\n", - " ax[m][0].plot(\n", - " range(20),\n", - " 627.5 * es_e,\n", - " marker=\"^\",\n", - " label=\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\",\n", - " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", - " linewidth=3.5,\n", - " markersize=14,\n", - " linestyle=\"--\",\n", - " )\n", - " ax[m][1].plot(\n", - " range(20),\n", - " 627.5 * (es_e - ref_e),\n", - " marker=\"^\",\n", - " label=\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\",\n", - " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", - " linewidth=3.5,\n", - " markersize=14,\n", - " linestyle=\"--\",\n", - " )\n", - " ax[m][0].errorbar(\n", - " x=range(20),\n", - " y=627.5 * ref_e,\n", - " yerr=627.5 * std,\n", - " label=\"Reference deep QMC\",\n", - " color=\"k\",\n", - " linestyle=\"-\",\n", - " linewidth=3,\n", - " alpha=0.8,\n", - " )\n", - " ax[m][1].axhline(-1, color=\"k\", ls=\":\", linewidth=3)\n", - " ax[m][1].axhline(1, color=\"k\", ls=\":\", linewidth=3)\n", - " if m != 2:\n", - " ax[m][1].text(\n", - " 0.56,\n", - " 0.9,\n", - " f\"MARE={627.5*np.mean(np.abs((qmc_e-ref_e))):.2f}\",\n", + "def make_reaction_profile(fig, ax):\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", + " qmc_e = raw_energy.loc[\n", + " (raw_energy[\"MEP_system\"] == j)\n", + " & (raw_energy[\"Ansatz\"] == \"Orbformer fine-tune LAC (400k), all structures\")\n", + " & (raw_energy[\"Fine-tune/Train_steps\"] == step)\n", + " ][\"E(Hartree)\"]\n", + " qmc_e = relative(np.array(qmc_e))\n", + " es_e = es_energies.loc[\n", + " (es_energies[\"MEP_system\"] == j)\n", + " & (es_energies[\"Method\"] == \"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\")\n", + " ][\"E(Hartree)\"]\n", + " es_e = relative(np.array(es_e))\n", + " ref_e = relative(np.array(ref.loc[ref[\"MEP_system\"] == j][\"Reference_E(Hartree)\"]))\n", + " std = np.array(ref.loc[ref[\"MEP_system\"] == j][\"Reference_E_std(Hartree)\"])\n", + " ax[m][0].plot(\n", + " range(20),\n", + " 627.5 * qmc_e,\n", " color=\"#27474a\",\n", - " transform=ax[m][1].transAxes,\n", - " fontsize=22,\n", + " marker=\"o\",\n", + " label=\"Best Orbformer fine-tune LAC (400k), all structures\",\n", + " linewidth=3.5,\n", + " markersize=14,\n", + " linestyle=\"-.\",\n", + " markeredgecolor=(1, 1, 1, 0.5),\n", + " markeredgewidth=2,\n", " )\n", - " ax[m][1].text(\n", - " 0.56,\n", - " 0.8,\n", - " f\"MARE={627.5*np.mean(np.abs((es_e-ref_e))):.2f}\",\n", - " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", - " transform=ax[m][1].transAxes,\n", - " fontsize=22,\n", - " )\n", - " else:\n", - " ax[m][1].text(\n", - " 0.56,\n", - " 0.8,\n", - " f\"MARE={627.5*np.mean(np.abs((qmc_e-ref_e))):.2f}\",\n", + " ax[m][1].plot(\n", + " range(20),\n", + " 627.5 * (qmc_e - ref_e),\n", " color=\"#27474a\",\n", - " transform=ax[m][1].transAxes,\n", - " fontsize=22,\n", + " marker=\"o\",\n", + " label=\"Best Orbformer fine-tune LAC (400k), all structures\",\n", + " linewidth=3.5,\n", + " markersize=14,\n", + " linestyle=\"-.\",\n", + " markeredgecolor=(1, 1, 1, 0.5),\n", + " markeredgewidth=2,\n", " )\n", - " ax[m][1].text(\n", - " 0.56,\n", - " 0.7,\n", - " f\"MARE={627.5*np.mean(np.abs((es_e-ref_e))):.2f}\",\n", + " ax[m][0].plot(\n", + " range(20),\n", + " 627.5 * es_e,\n", + " marker=\"^\",\n", + " label=\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\",\n", " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", - " transform=ax[m][1].transAxes,\n", - " fontsize=22,\n", + " linewidth=3.5,\n", + " markersize=14,\n", + " linestyle=\"--\",\n", + " markeredgecolor=(1, 1, 1, 0.5),\n", + " markeredgewidth=2,\n", " )\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[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", - " 0.03, 0.9, \"(c) 2-Aminopropan-2-ol\", transform=ax[2][0].transAxes, fontsize=24\n", - ")\n", - "ax[2][1].text(\n", - " 0.03, 0.9, \"(f) 2-Aminopropan-2-ol\", transform=ax[2][1].transAxes, fontsize=24\n", - ")\n", - "\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_0.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.1)\n", - "ab = AnnotationBbox(imagebox, (1.1, -55), frameon=False)\n", - "ax[0][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_8.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.2)\n", - "ab = AnnotationBbox(imagebox, (7.6, 2), frameon=False)\n", - "ax[0][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_19.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.18)\n", - "ab = AnnotationBbox(imagebox, (17, -20), frameon=False)\n", - "ax[0][0].add_artist(ab)\n", - "ax[0][0].annotate(\"\", xytext=(0, -85), xy=(0.7, -70), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[0][0].annotate(\"\", xytext=(8, 45), xy=(8, 10), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[0][0].annotate(\"\", xytext=(19, 10), xy=(18, -10), arrowprops=dict(arrowstyle=\"->\"))\n", - "\n", + " ax[m][1].plot(\n", + " range(20),\n", + " 627.5 * (es_e - ref_e),\n", + " marker=\"^\",\n", + " label=\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\",\n", + " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", + " linewidth=3.5,\n", + " markersize=14,\n", + " linestyle=\"--\",\n", + " markeredgecolor=(1, 1, 1, 0.5),\n", + " markeredgewidth=2,\n", + " )\n", + " ax[m][0].errorbar(\n", + " x=range(20),\n", + " y=627.5 * ref_e,\n", + " yerr=627.5 * std,\n", + " label=\"Reference deep QMC\",\n", + " color=\"k\",\n", + " linestyle=\"-\",\n", + " linewidth=3,\n", + " alpha=0.8,\n", + " )\n", + " ax[m][1].axhline(-1, color=\"k\", ls=\":\", linewidth=3, zorder=0)\n", + " ax[m][1].axhline(1, color=\"k\", ls=\":\", linewidth=3, zorder=0)\n", + " ax[m][1].set_ylim(\n", + " min(627.5 * (es_e - ref_e)) - 1, max(627.5 * (es_e - ref_e)) + [1.5, 3, 2][m]\n", + " )\n", + " if m != 2:\n", + " ax[m][1].text(\n", + " 0.56,\n", + " 0.9,\n", + " f\"MARE={627.5*np.mean(np.abs((qmc_e-ref_e))):.2f}\",\n", + " color=\"#27474a\",\n", + " transform=ax[m][1].transAxes,\n", + " fontsize=22,\n", + " )\n", + " ax[m][1].text(\n", + " 0.56,\n", + " 0.8,\n", + " f\"MARE={627.5*np.mean(np.abs((es_e-ref_e))):.2f}\",\n", + " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", + " transform=ax[m][1].transAxes,\n", + " fontsize=22,\n", + " )\n", + " else:\n", + " ax[m][1].text(\n", + " 0.56,\n", + " 0.8,\n", + " f\"MARE={627.5*np.mean(np.abs((qmc_e-ref_e))):.2f}\",\n", + " color=\"#27474a\",\n", + " transform=ax[m][1].transAxes,\n", + " fontsize=22,\n", + " )\n", + " ax[m][1].text(\n", + " 0.56,\n", + " 0.7,\n", + " f\"MARE={627.5*np.mean(np.abs((es_e-ref_e))):.2f}\",\n", + " color=colormaps_es[\"DF-MRCISD+Q(10,10)/aug-cc-pVDZ\"],\n", + " transform=ax[m][1].transAxes,\n", + " fontsize=22,\n", + " )\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[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(0.03, 0.9, \"(c) 2-Aminopropan-2-ol\", transform=ax[2][0].transAxes, fontsize=24)\n", + " ax[2][1].text(0.03, 0.9, \"(f) 2-Aminopropan-2-ol\", transform=ax[2][1].transAxes, fontsize=24)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_0.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.1)\n", + " ab = AnnotationBbox(imagebox, (1.1, -55), frameon=False)\n", + " ax[0][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_8.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.2)\n", + " ab = AnnotationBbox(imagebox, (7.6, 2), frameon=False)\n", + " ax[0][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Ethane/Ethane_19.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.18)\n", + " ab = AnnotationBbox(imagebox, (17, -20), frameon=False)\n", + " ax[0][0].add_artist(ab)\n", + " ax[0][0].annotate(\"\", xytext=(0, -85), xy=(0.7, -70), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[0][0].annotate(\"\", xytext=(8, 45), xy=(8, 10), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[0][0].annotate(\"\", xytext=(19, 10), xy=(18, -10), arrowprops=dict(arrowstyle=\"->\"))\n", "\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_0.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.1)\n", - "ab = AnnotationBbox(imagebox, (1.3, -40), frameon=False)\n", - "ax[1][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_6.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.18)\n", - "ab = AnnotationBbox(imagebox, (8.8, 40), frameon=False)\n", - "ax[1][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_10.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.18)\n", - "ab = AnnotationBbox(imagebox, (9.2, -20), frameon=False)\n", - "ax[1][0].add_artist(ab)\n", - "arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_19.png\")\n", - "imagebox = OffsetImage(arrimg, zoom=0.2)\n", - "ab = AnnotationBbox(imagebox, (17, 30), frameon=False)\n", - "ax[1][0].add_artist(ab)\n", - "ax[1][0].annotate(\"\", xytext=(0, -65), xy=(0.7, -50), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[1][0].annotate(\"\", xytext=(6.4, 45), xy=(7.5, 45), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[1][0].annotate(\"\", xytext=(10, -42), xy=(10, -30), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[1][0].annotate(\"\", xytext=(19, 55), xy=(18, 40), arrowprops=dict(arrowstyle=\"->\"))\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_0.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.1)\n", + " ab = AnnotationBbox(imagebox, (1.3, -40), frameon=False)\n", + " ax[1][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_6.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.18)\n", + " ab = AnnotationBbox(imagebox, (8.8, 40), frameon=False)\n", + " ax[1][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_10.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.18)\n", + " ab = AnnotationBbox(imagebox, (9.2, -20), frameon=False)\n", + " ax[1][0].add_artist(ab)\n", + " arrimg = mpimg.imread(f\"{data_dir}/molecule_images/bbmep/Formamide/Formamide_19.png\")\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.2)\n", + " ab = AnnotationBbox(imagebox, (17, 30), frameon=False)\n", + " ax[1][0].add_artist(ab)\n", + " ax[1][0].annotate(\"\", xytext=(0, -65), xy=(0.7, -50), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[1][0].annotate(\"\", xytext=(6.4, 45), xy=(7.5, 45), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[1][0].annotate(\"\", xytext=(10, -42), xy=(10, -30), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[1][0].annotate(\"\", xytext=(19, 55), xy=(18, 40), arrowprops=dict(arrowstyle=\"->\"))\n", "\n", - "arrimg = mpimg.imread(\n", - " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_0.png\"\n", - ")\n", - "imagebox = OffsetImage(arrimg, zoom=0.18)\n", - "ab = AnnotationBbox(imagebox, (1.4, -25), frameon=False)\n", - "ax[2][0].add_artist(ab)\n", - "arrimg = mpimg.imread(\n", - " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_19.png\"\n", - ")\n", - "arrimg = remove_black_edge(rotate(arrimg, angle=60))\n", - "imagebox = OffsetImage(arrimg, zoom=0.22)\n", - "ab = AnnotationBbox(imagebox, (17, 0), frameon=False)\n", - "ax[2][0].add_artist(ab)\n", - "ax[2][0].annotate(\"\", xytext=(0, -50), xy=(0.7, -40), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[2][0].annotate(\"\", xytext=(19, 35), xy=(17, 20), arrowprops=dict(arrowstyle=\"->\"))\n", - "ax[0][0].text(0.28, 1.1, \"Relative Energy\", transform=ax[0][0].transAxes, fontsize=26)\n", - "ax[0][1].text(\n", - " 0.18, 1.1, \"Relative Energy Error\", transform=ax[0][1].transAxes, fontsize=26\n", - ")\n", - "ax[2][0].set_ylim([-65, 58])\n", - "plt.subplots_adjust(wspace=0.2, hspace=0.02)\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", + " arrimg = mpimg.imread(\n", + " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_0.png\"\n", + " )\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.18)\n", + " ab = AnnotationBbox(imagebox, (1.4, -25), frameon=False)\n", + " ax[2][0].add_artist(ab)\n", + " arrimg = mpimg.imread(\n", + " f\"{data_dir}/molecule_images/bbmep/2-Aminopropan-2-ol/2-Aminopropan-2-ol_19.png\"\n", + " )\n", + " arrimg = remove_black_edge(rotate(arrimg, angle=60))\n", + " imagebox = OffsetImage(remove_background(arrimg), zoom=0.22)\n", + " ab = AnnotationBbox(imagebox, (17, 0), frameon=False)\n", + " ax[2][0].add_artist(ab)\n", + " ax[2][0].annotate(\"\", xytext=(0, -50), xy=(0.7, -40), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[2][0].annotate(\"\", xytext=(19, 35), xy=(17, 20), arrowprops=dict(arrowstyle=\"->\"))\n", + " ax[0][0].text(0.28, 1.03, \"Relative Energy\", transform=ax[0][0].transAxes, fontsize=26)\n", + " ax[0][1].text(0.18, 1.03, \"Relative Energy Error\", transform=ax[0][1].transAxes, fontsize=26)\n", + " ax[0][0].set_ylim([-99, 74])\n", + " ax[2][0].set_ylim([-65, 58])\n", + " plt.subplots_adjust(wspace=0.2, hspace=0.02)\n", + " for i in range(6):\n", + " ax[i % 3][i // 3].set_xticks(5 * np.arange(4))\n", + " ax[i % 3][i // 3].set_xticklabels([])\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(5 * np.arange(4))\n", + " ax[2][0].set_xlabel(\"Image ID on MEP\", fontsize=22)\n", + " ax[2][1].set_xticklabels(5 * np.arange(4))\n", + " ax[2][1].set_xlabel(\"Image ID on MEP\", fontsize=22)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c58b9197", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(3, 2, figsize=(13, 17.5))\n", + "make_reaction_profile(fig, ax)\n", "plt.draw()\n", "if save_figures:\n", - " plt.savefig(f\"{data_dir}/BBMEP_energy_profile.pdf\", dpi=600, bbox_inches='tight')\n", + " plt.savefig(f\"{data_dir}/BBMEP_energy_profile.pdf\", dpi=600, bbox_inches=\"tight\")\n", "plt.show()" ] }, @@ -551,12 +613,33 @@ "id": "e4c631ee-d7aa-4571-b2a0-7b72e4726ec1", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "fig, ax = plt.subplots(\n", + " 3,\n", + " 7,\n", + " figsize=(31.5, 19),\n", + " gridspec_kw={\n", + " \"width_ratios\": [1, 0.2, 1, 0.25, 1.5, 0.14, 1.5],\n", + " \"height_ratios\": [1, 1, 1],\n", + " \"wspace\": 0,\n", + " \"hspace\": 0.01,\n", + " },\n", + ")\n", + "for empty in ax[:, 1::2]:\n", + " for e in empty:\n", + " e.axis(\"off\")\n", + "make_reaction_profile(fig, ax[:, :4:2])\n", + "make_scaling(fig, ax[:, 4::2], legend_offset=(0.045, 0.12), mols_offset=(-0.02, -0.1))\n", + "plt.subplots_adjust(left=0.15, right=0.99, bottom=0.2, top=0.93)\n", + "if save_figures:\n", + " plt.savefig(f\"{data_dir}/BBMEP_combined.pdf\", dpi=600, bbox_inches=\"tight\")\n", + "plt.show()" + ] } ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -570,7 +653,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.10" + "version": "3.11.13" } }, "nbformat": 4,