00_main_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9d5bcaa-3292-44f9-ab96-04db5f58faeb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%config Completer.use_jedi = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01476047-0e63-40c1-ac8f-3efac512b4dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import json\n",
    "from pathlib import Path\n",
    "import string\n",
    "import sys\n",
    "\n",
    "from rdkit import Chem\n",
    "from IPython.display import clear_output\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.spatial import distance_matrix\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import tqdm\n",
    "import torch\n",
    "from sklearn.metrics import balanced_accuracy_score, accuracy_score, f1_score, recall_score\n",
    "\n",
    "from multimodal_molecules.core import Ensemble, Estimator, get_data\n",
    "from multimodal_molecules.plotting import set_defaults, set_grids, density_scatter, remove_axis_spines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3d6b099-8813-4492-b487-e1baf35a1434",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_defaults()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c116664f-6fa4-4b8e-872f-72e50e3c24ee",
   "metadata": {},
   "source": [
    "# Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06081382-d809-46cb-9939-0811cce923b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_deviating_estimators(preds):\n",
    "    N = preds.shape[0]\n",
    "    assert N % 2 == 0\n",
    "    s = preds.round().sum(axis=0)\n",
    "    where = np.where(s > N//2)\n",
    "    s[where] = N - s[where]\n",
    "    return s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355720cb-44f0-4a90-a8c9-b100e523dfb8",
   "metadata": {},
   "source": [
    "# Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11c70a18-2fd7-44ba-846c-3dfad124541b",
   "metadata": {},
   "outputs": [],
   "source": [
    "C_grid = np.loadtxt(\"data/c_grid.txt\")\n",
    "N_grid = np.loadtxt(\"data/n_grid.txt\")\n",
    "O_grid = np.loadtxt(\"data/o_grid.txt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19e36cd8-39e7-4b51-bd19-14965d183a1c",
   "metadata": {},
   "source": [
    "# Multimodal PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5612be4a-1c86-4346-87b0-577a428cead9",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = get_data(elements=\"CNO\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fab1cd7c-b051-4318-858c-ca2977756112",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca_CNO = PCA(2)\n",
    "scaler = StandardScaler()\n",
    "pca_CNO.fit(scaler.fit_transform(data[\"X_test\"]))\n",
    "# pca_CNO.fit(data[\"X_test\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b987ad11-8442-4f94-8468-ba0842c3001b",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_CNO = pca_CNO.transform(scaler.transform(data[\"X_test\"]))\n",
    "# w_CNO = pca_CNO.transform(data[\"X_test\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a81ec6b-e145-48c4-9ad9-16db0aedb0b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_example(p, w=w_CNO, keep=10):\n",
    "    p = np.array(p)\n",
    "    delta = np.abs(w - p).sum(axis=1)\n",
    "    argsorted = np.argsort(delta)\n",
    "    return argsorted[:keep]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "871aa05f-73b7-4428-a90b-b729202ee08a",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = mpl.colormaps[\"viridis\"].resampled(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75a0015a-b183-45dd-8dac-c99ef14128e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Edge points\n",
    "p1 = get_example([-2, 29])\n",
    "p2 = get_example([20, 5])\n",
    "p3 = get_example([5, -20])\n",
    "p4 = get_example([-22, -3])\n",
    "p5 = get_example([-16, 5])\n",
    "p6 = get_example([-1.5, 1])\n",
    "p7 = get_example([-2.5, -5])\n",
    "p8 = get_example([12.5, -9.5])\n",
    "selected_points = [p1, p2, p3, p4, p5, p6, p7, p8]\n",
    "colors = [cmap(ii) for ii in range(len(selected_points))]\n",
    "smiles = [[data[\"smiles_test\"][xx] for xx in yy] for yy in selected_points]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf4c0fa-38e5-451d-833f-4b3ef4394bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(2, 2), sharey=True)\n",
    "\n",
    "# remove_axis_spines(ax, visible=True)\n",
    "\n",
    "density_scatter(w_CNO[:, 0], w_CNO[:, 1], ax, bins=20, s=1, cmap=\"binary\", rasterized=True)\n",
    "\n",
    "for point, color in zip(selected_points, colors):\n",
    "    ax.scatter(w_CNO[point[0], 0], w_CNO[point[0], 1], color=color, marker=\"o\", s=20, facecolors='none', linewidth=1)\n",
    "\n",
    "ax.set_xlabel(r\"$z_1$ [a.u.]\")\n",
    "ax.set_ylabel(r\"$z_2$ [a.u.]\")\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.text(0.9, 0.9, \"(a)\", ha=\"center\", va=\"center\", transform=ax.transAxes)\n",
    "\n",
    "\n",
    "plt.show()\n",
    "# plt.savefig(\"figures/fig_pca_1/fig_pca_1.svg\", bbox_inches=\"tight\", dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4eed256-981a-44af-bf86-e9cebe59e915",
   "metadata": {},
   "outputs": [],
   "source": [
    "letters = [\"i\", \"ii\", \"iii\", \"iv\", \"v\", \"vi\", \"vii\", \"viii\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41d5a282-96e9-4f80-b3d3-119080bf7fad",
   "metadata": {},
   "outputs": [],
   "source": [
    "lw = 0.5\n",
    "offset = 1\n",
    "C_limit = 160\n",
    "N_limit = 160\n",
    "O_limit = 120\n",
    "\n",
    "fig, all_axs = plt.subplots(8, 3, figsize=(5, 8), sharey=True)\n",
    "\n",
    "for ii, (points, smile, color, letter, axs) in enumerate(zip(selected_points, smiles, colors, letters, all_axs)):\n",
    "    \n",
    "    # fig, axs = plt.subplots(1, 3, figsize=(6, 1), sharey=True)\n",
    "\n",
    "    axs[0].text(0.1, 0.9, f\"({letter})\", ha=\"center\", va=\"center\", transform=axs[0].transAxes)\n",
    "    \n",
    "    for ax in axs:\n",
    "\n",
    "        for p in points:\n",
    "            ax.plot(C_grid[:C_limit], data[\"X_test\"][p, :C_limit] + offset, color=color, linewidth=lw)\n",
    "        \n",
    "        for p in points:\n",
    "            ax.plot(N_grid[:N_limit], data[\"X_test\"][p, 200:200+N_limit] + offset, color=color, linewidth=lw)\n",
    "        \n",
    "        for p in points:\n",
    "            ax.plot(O_grid[:O_limit], data[\"X_test\"][p, 400:400+O_limit] + offset, color=color, linewidth=lw)\n",
    "\n",
    "        ax.set_yticks([])\n",
    "        ax.spines[[\"right\", \"top\"]].set_visible(False)\n",
    "        ax.tick_params(\n",
    "            which='both', direction='in', bottom=True, left=True\n",
    "        )\n",
    "\n",
    "    axs[0].set_xlim(C_grid[0], C_grid[C_limit])\n",
    "    axs[0].set_xticks([280, 315])\n",
    "    axs[1].set_xlim(N_grid[0], N_grid[N_limit])\n",
    "    axs[1].set_xticks([400, 432])\n",
    "    axs[2].set_xlim(O_grid[0], O_grid[O_limit])\n",
    "    axs[2].set_xticks([533, 558])\n",
    "\n",
    "    # Hide specific spines\n",
    "    axs[1].spines.left.set_visible(False)\n",
    "    axs[2].spines.left.set_visible(False)\n",
    "\n",
    "    if ii == len(all_axs) - 1 or ii == len(all_axs) // 2 - 1:\n",
    "        axs[0].set_ylabel(r\"$\\mu(E)$ [a.u.]\")\n",
    "        axs[1].set_xlabel(r\"$E$ [e.V.]\")\n",
    "    else:\n",
    "        for ax in axs:\n",
    "            ax.set_xticklabels([])\n",
    "\n",
    "    d = 1.5  # proportion of vertical to horizontal extent of the slanted line\n",
    "    kwargs = dict(marker=[(-1, -d), (1, d)], markersize=4, linestyle=\"none\", color='k', mec='k', mew=0.75, clip_on=False)\n",
    "    axs[0].plot([1, 1], [0, 0], transform=axs[0].transAxes, **kwargs)\n",
    "    axs[1].plot([0, 0], [0, 0], transform=axs[1].transAxes, **kwargs)\n",
    "    axs[1].plot([1, 1], [0, 0], transform=axs[1].transAxes, **kwargs)\n",
    "    axs[2].plot([0, 0], [0, 0], transform=axs[2].transAxes, **kwargs)\n",
    "        \n",
    "    # Get the smiles...\n",
    "    # svgs = Chem.Draw.MolsToGridImage([Chem.MolFromSmiles(s) for s in smile], useSVG=True)\n",
    "    # with open(f\"figures/fig_pca_1/example_structures_{letter}.svg\", \"w\") as f:\n",
    "    #     f.write(svgs.data)\n",
    "\n",
    "plt.subplots_adjust(wspace=0.05, hspace=0.2)\n",
    "# plt.savefig(\"figures/fig_pca_1/example_spectra.svg\", bbox_inches=\"tight\", dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b29bdcf-53e2-424c-9e26-4252bcf3d8f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "for xx in indexes[2:]:\n",
    "    smiles = [data[\"smiles_test\"][ii] for ii in xx]\n",
    "    svgs = Chem.Draw.MolsToGridImage([Chem.MolFromSmiles(s) for s in smiles], useSVG=True)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36f20cc4-b42a-4db0-8c5f-4d5a3198f3cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "svgs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68307f2b-1c7c-437a-bd90-3f7f5eca4fc0",
   "metadata": {},
   "source": [
    "# Carbon exploration and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a240591-1655-44e5-810e-b77632d3aaaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ELEMENT = \"C\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92cd744e-8caf-4274-9bd7-f20875031ce3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble = Ensemble.from_path(f\"data/23-12-06_torch_models/{ELEMENT}\")\n",
    "data = get_data(elements=ELEMENT)\n",
    "grid = np.loadtxt(f\"data/{ELEMENT.lower()}_grid.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce9e38d-fda1-424e-9728-a46ff9012e70",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_spectrum = data[\"X_train\"].mean(axis=0).squeeze()\n",
    "for ii, (e, s) in enumerate(zip(grid, mean_spectrum)):\n",
    "    if s > 1:\n",
    "        break\n",
    "print(grid[ii])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf01d19-894f-4f66-8f24-863f6dfab72e",
   "metadata": {},
   "source": [
    "## PCA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f33d0c46-61e5-480b-a83b-a3619b56a500",
   "metadata": {},
   "source": [
    "We make some nice plots here by doing the following:\n",
    "1. Decompose the spectral data into 2 dimensions\n",
    "2. Select interesting functional groups\n",
    "3. Plot the first two principal components of the data against labeled data\n",
    "4. Plot using 2d density histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5982b84c-3868-40be-ac3b-8bb6e679b921",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(2)\n",
    "X = data[\"X_test\"]\n",
    "w = pca.fit_transform(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caafe07a-4499-4cf6-8fec-6dcf575bb432",
   "metadata": {},
   "outputs": [],
   "source": [
    "interesting_functional_groups = [\n",
    "    # None,\n",
    "    \"Quaternary_carbon\",\n",
    "    \"Alkene\",\n",
    "    \"Alkyne\",\n",
    "    \"Aromatic\",\n",
    "]\n",
    "N = len(interesting_functional_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f3f209-e4a7-4095-aaac-fde66a115c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "smiles = [\"CC(C)(C)(C)\", \"C(C)(C)=C(C)(C)\", \"C(C)#C(C)\", \"c1ccccc1\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c087d3e-3193-4eaf-a441-0a77ad20fa9d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "svgs = Chem.Draw.MolsToGridImage([Chem.MolFromSmiles(s) for s in smiles], useSVG=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e04ce3-c2cd-4ab4-8d0c-8e487aa0e7a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "svgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de2598c6-211c-486f-8c6d-c16308ba74ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with open(\"figures/fig_pca_2-carbon/motifs.svg\", \"w\") as f:\n",
    "    f.write(svgs.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace49cb4-5573-425a-b907-c5ed4491a355",
   "metadata": {},
   "outputs": [],
   "source": [
    "chars = [\"i\", \"ii\", \"iii\", \"iv\"]\n",
    "fig, axs = plt.subplots(1, N, figsize=(2 * N, 2), sharey=True)\n",
    "\n",
    "for fg, ax, label in zip(interesting_functional_groups, axs, chars):\n",
    "    remove_axis_spines(ax, visible=True)\n",
    "\n",
    "    density_scatter(w[:, 0], w[:, 1], ax, bins=20, s=0.2, cmap=\"binary\", rasterized=True)\n",
    "    \n",
    "    if fg is not None:\n",
    "        index = data[\"functional_groups\"].index(fg)\n",
    "        condition = (data[\"Y_test\"][:, index] == 1)\n",
    "        where = np.where(condition)[0]\n",
    "        density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "\n",
    "    ax.text(0.05, 0.95, f\"({label})\", ha=\"left\", va=\"top\", transform=ax.transAxes)\n",
    "\n",
    "axs[0].set_xlabel(r\"$z_1$~[a.u.]\")\n",
    "axs[0].set_ylabel(r\"$z_2$~[a.u.]\")\n",
    "\n",
    "plt.subplots_adjust(wspace=0.1)\n",
    "# plt.savefig(\"figures/fig_pca_2-carbon/pca.svg\", bbox_inches='tight', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125fda28-6942-422a-86d6-4a8dc427fdcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "mult = 2\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(1.5, 2), sharey=True)\n",
    "\n",
    "# remove_axis_spines(ax, visible=True)\n",
    "\n",
    "cmap = mpl.colormaps[\"rainbow\"].resampled(len(interesting_functional_groups))\n",
    "\n",
    "for ii, (fg, label) in enumerate(zip(interesting_functional_groups, chars)):\n",
    "    \n",
    "    index = data[\"functional_groups\"].index(fg)\n",
    "    condition = (data[\"Y_test\"][:, index] == 1)\n",
    "    where = np.where(condition)[0]\n",
    "    X_where = X[where, :]\n",
    "    mu = X_where.mean(axis=0)\n",
    "    sd = X_where.std(axis=0)\n",
    "    ax.plot(grid, mu - ii*mult, color=cmap(ii))\n",
    "    for jj in range(10):\n",
    "        ax.plot(grid, X_where[jj, :] - ii * mult, color=cmap(ii), alpha = 0.5, linewidth=0.5)\n",
    "    # ax.fill_between(grid, mu + ii*mult - sd, mu + ii*mult + sd, linewidth=0, color=cmap(ii), alpha=0.2)\n",
    "\n",
    "# ax.legend()\n",
    "ax.set_yticks([])\n",
    "set_grids(ax, grid=False)\n",
    "ax.set_ylabel(r\"$\\mu(E)$~[a.u.]\")\n",
    "ax.set_xlabel(r\"$E$~[eV]\")\n",
    "ax.set_xticks([280, 290, 300, 310])\n",
    "ax.set_xlim([275, 315])\n",
    "\n",
    "# plt.subplots_adjust(wspace=0.1)\n",
    "# plt.savefig(\"figures/fig_pca_2-carbon/spectra.svg\", bbox_inches='tight', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af55517-25ef-44da-8bc6-231939297b16",
   "metadata": {},
   "source": [
    "# Nitrogen exploration and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0947fd4d-d8ee-48f6-bc29-ab8349119f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ELEMENT = \"N\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aecd7a5-bba2-47ec-bcef-3f65a8571423",
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble = Ensemble.from_path(f\"data/23-12-06_torch_models/{ELEMENT}\")\n",
    "data = get_data(elements=ELEMENT)\n",
    "grid = np.loadtxt(f\"data/{ELEMENT.lower()}_grid.txt\")\n",
    "X = data[\"X_test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7935f8-5ace-436c-aa9b-69f0158c3d10",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_spectrum = data[\"X_train\"].mean(axis=0).squeeze()\n",
    "for ii, (e, s) in enumerate(zip(grid, mean_spectrum)):\n",
    "    if s > 1:\n",
    "        break\n",
    "print(grid[ii])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f3d3af2-0ad1-4852-985c-b30d76147776",
   "metadata": {},
   "source": [
    "## PCA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "815389f7-c321-45ea-b091-8b4b6f2a96eb",
   "metadata": {},
   "source": [
    "We make some nice plots here by doing the following:\n",
    "1. Decompose the spectral data into 2 dimensions\n",
    "2. Select interesting functional groups\n",
    "3. Plot the first two principal components of the data against labeled data\n",
    "4. Plot using 2d density histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40a9ddb0-2aa4-4553-a62d-4d9ca177ab5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(2)\n",
    "w = pca.fit_transform(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9534c90b-8d0c-423b-9df2-1edbca989435",
   "metadata": {},
   "outputs": [],
   "source": [
    "data[\"functional_groups\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a50e8b77-308b-4db4-bf80-f2bdf57658a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "interesting_functional_groups = [\n",
    "    # None,\n",
    "    \"Amine\",\n",
    "    \"Amide\",\n",
    "    \"Nitrile\",\n",
    "    \"Heteroaromatic\",\n",
    "    # \"Secondary_aliph_amine\",\n",
    "    # \"Tertiary_aliph_amine\",\n",
    "]\n",
    "N = len(interesting_functional_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa55e73-f510-4ed6-b12a-003db359e793",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, N, figsize=(2 * N, 2), sharey=True)\n",
    "\n",
    "for ax in axs:\n",
    "    ax.set_aspect(\"equal\")\n",
    "\n",
    "for fg, ax in zip(interesting_functional_groups, axs):\n",
    "    remove_axis_spines(ax, visible=True)\n",
    "\n",
    "    density_scatter(w[:, 0], w[:, 1], ax, bins=20, s=0.2, cmap=\"Blues\", rasterized=True)\n",
    "    \n",
    "    if fg is not None:\n",
    "        index = data[\"functional_groups\"].index(fg)\n",
    "        condition = (data[\"Y_test\"][:, index] == 1)\n",
    "\n",
    "        # Plot where one nitrogen\n",
    "        condition_only_one = [Counter(xx.lower())[\"n\"] == 1 for xx in data[\"smiles_test\"]]\n",
    "        where = np.where(condition & condition_only_one)[0]\n",
    "        density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "\n",
    "        # Plot where two nitrogen\n",
    "        # condition_only_one = [Counter(xx.lower())[\"n\"] == 2 for xx in data[\"smiles_test\"]]\n",
    "        # where = np.where(condition & condition_only_one)[0]\n",
    "        # density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "        \n",
    "\n",
    "    # ax.text(0.05, 0.95, f\"({label})\", ha=\"left\", va=\"top\", transform=ax.transAxes)\n",
    "\n",
    "axs[0].set_xlabel(r\"$z_1$~[a.u.]\")\n",
    "axs[0].set_ylabel(r\"$z_2$~[a.u.]\")\n",
    "\n",
    "plt.subplots_adjust(wspace=0.1)\n",
    "plt.show()\n",
    "# plt.savefig(\"figures/fig_pca_2-nitrogen/pca.svg\", bbox_inches=\"tight\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3861f070-a0fa-43be-b121-c84778e3d421",
   "metadata": {},
   "source": [
    "### Plots of actual spectra examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "481c6912-b38a-4000-8fa3-72dde52c6b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "mult = 4\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2, 2), sharey=True)\n",
    "\n",
    "# remove_axis_spines(ax, visible=True)\n",
    "\n",
    "cmap = mpl.colormaps[\"rainbow\"].resampled(len(interesting_functional_groups))\n",
    "\n",
    "for ii, (fg,) in enumerate(zip(interesting_functional_groups)):\n",
    "    \n",
    "    index = data[\"functional_groups\"].index(fg)\n",
    "    condition = (data[\"Y_test\"][:, index] == 1)\n",
    "    condition_only_one = [Counter(xx.lower())[\"n\"] == 1 for xx in data[\"smiles_test\"]]\n",
    "    \n",
    "    where = np.where(condition & condition_only_one)[0]\n",
    "    \n",
    "    X_where = X[where, :]\n",
    "    mu = X_where.mean(axis=0)\n",
    "    sd = X_where.std(axis=0)\n",
    "    ax.plot(grid, mu - ii*mult, color=cmap(ii), zorder=-ii)\n",
    "    for jj in range(100, 110):\n",
    "        ax.plot(grid, X_where[jj, :] - ii * mult, color=cmap(ii), alpha = 0.5, linewidth=0.5)\n",
    "    # ax.fill_between(grid, mu + ii*mult - sd, mu + ii*mult + sd, linewidth=0, color=cmap(ii), alpha=0.2)\n",
    "\n",
    "ax.set_yticks([])\n",
    "set_grids(ax, grid=False)\n",
    "ax.set_ylabel(r\"$\\mu(E)$~[a.u.]\")\n",
    "ax.set_xlabel(r\"$E$~[eV]\")\n",
    "ax.set_xticks([400, 410, 420, 430])\n",
    "ax.set_xlim([397, 431])\n",
    "\n",
    "# plt.subplots_adjust(wspace=0.1)\n",
    "plt.savefig(\"figures/fig_pca_3-nitrogen-oxygen/spectra.svg\", bbox_inches='tight', dpi=300)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fca9a82-1683-49aa-88aa-14e3ba4db684",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Testing on experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c340b7b5-d712-4a12-a538-b055ea7736f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"experiment/n_exp.json\", \"r\") as f:\n",
    "    exp_data = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db8d23c2-2688-4265-bcf5-331a3a4aeb8b",
   "metadata": {},
   "source": [
    "Get the infilled spectra from the loaded dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98e3d1c-f815-4e2c-8805-bcb0f8724d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "spectra_from_smiles = {\n",
    "    value[\"smiles\"]: value[\"infilled_spectrum\"]\n",
    "    for key, value in exp_data[\"data\"].items()\n",
    "    if key not in exp_data[\"errors\"]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c65019b2-2b41-4c9d-83c6-ce9b42cc5ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in spectra_from_smiles.items():\n",
    "    value = np.array(value)\n",
    "    value[value < 0] = 0\n",
    "    spectra_from_smiles[key] = value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e119d9c4-96e4-4a11-aa4e-0719a666e391",
   "metadata": {},
   "source": [
    "The index specifies the mapping between the smiles and targets. Note that we screen out the data we're not going to use on the functional group side of things here. For example, here we screen out any data that do not conform to the \"qm9\" standard (C, N, O, F, H)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c61305c-1280-4e6d-993f-1d66ffec2701",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _qm9_condition(smile):\n",
    "    if \"+\" in smile or \"-\" in smile:\n",
    "        return False\n",
    "    alphanumeric = [s for s in smile if s.isalpha()]\n",
    "    return all([xx.lower() in [\"c\", \"n\", \"o\", \"f\", \"h\"] for xx in alphanumeric])\n",
    "\n",
    "def _error_free_condition(smiles, s=spectra_from_smiles):\n",
    "    return smiles in s\n",
    "\n",
    "def get_index(\n",
    "    index=None,\n",
    "    require_error_free=True,\n",
    "    require_qm9=False,\n",
    "    require_functional_group=None,\n",
    "    feff_functional_groups=data[\"functional_groups\"]\n",
    "):\n",
    "    \n",
    "    if index is None:\n",
    "        index = pd.read_csv(\"experiment/index.csv\", index_col=0)\n",
    "\n",
    "    feff_mask = None\n",
    "    if feff_functional_groups is not None:\n",
    "\n",
    "        # We only want to keep the subset of experimental and simulation functional groups\n",
    "        common_fg = sorted(list(set(feff_functional_groups).intersection(set(index.columns[1:]))))\n",
    "\n",
    "        # Apply the mask\n",
    "        index = index[[\"SMILES\"] + common_fg]\n",
    "\n",
    "        # Compute the feff mask as a helper\n",
    "        feff_mask = np.array([ii for ii, fg in enumerate(feff_functional_groups) if fg in common_fg])\n",
    "        for mask_index, fg in zip(feff_mask, index.columns[1:]):\n",
    "            assert feff_functional_groups[mask_index] == fg\n",
    "\n",
    "    if require_qm9:\n",
    "        keep = index[\"SMILES\"].apply(_qm9_condition)\n",
    "        index = index[keep]\n",
    "\n",
    "    if require_error_free:\n",
    "        keep = index[\"SMILES\"].apply(_error_free_condition)\n",
    "        index = index[keep]\n",
    "\n",
    "    if require_functional_group is not None:\n",
    "        index = index[index[require_functional_group] == 1]\n",
    "\n",
    "    return index, feff_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49da314f-fe7b-4f05-9679-42c04e1cba30",
   "metadata": {},
   "outputs": [],
   "source": [
    "FUNCTIONAL_GROUP = \"Amine\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a999bca-e0da-4f2c-a541-7c491109baf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "index, feff_mask = get_index(require_functional_group=FUNCTIONAL_GROUP)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f368f4bd-09a0-45f8-9241-3da5f1c6dda4",
   "metadata": {},
   "source": [
    "Construct the new testing data (experimental results)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee6c3b58-0eef-409e-a048-4ef959e5d545",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_exp = []\n",
    "Y_exp = []\n",
    "for ii, row in index.iterrows():\n",
    "    smiles = row[\"SMILES\"]\n",
    "    Y_exp.append(row[1:].to_numpy())\n",
    "    X_exp.append(spectra_from_smiles[smiles])\n",
    "X_exp = np.array(X_exp)\n",
    "Y_exp = np.array(Y_exp, dtype=int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d036884f-a179-4aca-ade7-779fc90d0f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "feff_FG_index = data[\"functional_groups\"].index(FUNCTIONAL_GROUP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe11d07d-050b-4c86-b284-ceb84a4c6593",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = ensemble.predict(torch.FloatTensor(X_exp))\n",
    "mu = np.array(preds.round().mean(axis=0).round(), dtype=int)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f9157b-72bc-4434-9837-411baba39b15",
   "metadata": {},
   "source": [
    "This logic gets the number of deviating estimators, a proxy for uncertainty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bab06eb-5a2e-4b49-a5d4-b09546b32907",
   "metadata": {},
   "outputs": [],
   "source": [
    "deviating_estimators = get_deviating_estimators(preds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d028ba3e-e48d-4fbe-b081-efed7f454474",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Compare Amines from experiment and theory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0303a964-22cf-4b57-88d6-28813a6d6f82",
   "metadata": {},
   "source": [
    "First extract relevant information from FEFF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37961d5a-1e90-441f-b46b-a63df17ed4fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eeade6a-03e5-4981-a0e3-362bf1a8a846",
   "metadata": {},
   "outputs": [],
   "source": [
    "feff_smiles_test = data[\"smiles_test\"]\n",
    "where_one_n = np.array([1 if Counter(smile.lower())[\"n\"] == 1 else 0 for smile in feff_smiles_test])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a402567b-da10-4bb8-a1e5-4552fbfe624b",
   "metadata": {},
   "outputs": [],
   "source": [
    "locations = np.where(\n",
    "    (data[\"Y_test\"][:, feff_FG_index] == 1) & (where_one_n == 1)\n",
    ")[0] # Get all the Amines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f460ef-6e15-4a95-a508-978bc1a2d6a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "kept_feff_spectra = data[\"X_test\"][locations, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61dbcb0c-64b1-4a24-baae-94155c8d588d",
   "metadata": {},
   "outputs": [],
   "source": [
    "istar_exp = index.columns[1:].tolist()\n",
    "istar_exp = istar_exp.index(FUNCTIONAL_GROUP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bccdff0-7273-4de5-807d-0a089cef106a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_closest_spectrum(x1, X):\n",
    "    d = distance_matrix(x1.reshape(1, -1), X)\n",
    "    return np.argsort(d).squeeze()[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15157fd9-89ca-4f0f-8c69-461e1d768f7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_inputs = np.concatenate([data[\"X_train\"], data[\"X_val\"], data[\"X_test\"]], axis=0)\n",
    "all_smiles = data[\"smiles_train\"] + data[\"smiles_val\"] + data[\"smiles_test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed131092-99ca-4933-b78f-51a0c47a5110",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_canon_smiles = [Chem.CanonSmiles(smile) for smile in all_smiles]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c3cea3c-bfe8-4b28-a104-2f5d9410ef5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs=4\n",
    "\n",
    "fig, axs = plt.subplots(len(index), 1, figsize=(3, 2 * len(index)))\n",
    "\n",
    "cc = 0\n",
    "for _, row in index.iterrows():\n",
    "\n",
    "    ax = axs[cc]\n",
    "    \n",
    "    smiles = row[\"SMILES\"]\n",
    "    y = row[1:].to_numpy()\n",
    "    x = spectra_from_smiles[smiles]\n",
    "    pred = mu[cc, feff_FG_index].item()\n",
    "    dev = deviating_estimators[cc, feff_FG_index].item()\n",
    "    target = Y_exp[cc, istar_exp].item()\n",
    "    ax.plot(exp_data[\"feff_grid\"], x, color=\"black\", linewidth=0.5)\n",
    "\n",
    "    # See if smiles exists in training set\n",
    "    try:\n",
    "        s = Chem.CanonSmiles(smiles)\n",
    "        istar = all_canon_smiles.index(s)\n",
    "        x_star = all_inputs[istar]\n",
    "        ax.plot(exp_data[\"feff_grid\"], x_star, color=\"blue\", linewidth=0.5)\n",
    "        ax.text(0.99, 0.75, f\"smile in FEFF data: {s}\", ha=\"right\", va=\"top\", transform=ax.transAxes, fontsize=fs, color=\"blue\")\n",
    "    except ValueError:\n",
    "        pass\n",
    "\n",
    "    closest_x = get_closest_spectrum(np.array(x), data[\"X_test\"]).squeeze()\n",
    "    for ii in range(len(closest_x)):\n",
    "        closest_smiles = data[\"smiles_test\"][closest_x[ii]]\n",
    "        closest_smiles = closest_smiles.replace(\"#\", \"\\#\")\n",
    "        ax.plot(\n",
    "            exp_data[\"feff_grid\"], data[\"X_test\"][closest_x[ii], :],\n",
    "            color=\"red\", linewidth=0.5\n",
    "        )\n",
    "        ax.text(\n",
    "            0.99, 0.7 - ii*0.05, \"%s\" % closest_smiles,\n",
    "            ha=\"right\", va=\"top\", transform=ax.transAxes, fontsize=fs, color=\"red\"\n",
    "        )\n",
    "    # ax.plot(exp_data[\"feff_grid\"], kept_feff_spectra.mean(axis=0), color=\"red\")\n",
    "\n",
    "    ax.text(0.99, 0.95, f\"exp: {smiles}\", ha=\"right\", va=\"top\", transform=ax.transAxes, fontsize=fs)\n",
    "    ax.text(0.99, 0.9, f\"pred={pred}\", ha=\"right\", va=\"top\", transform=ax.transAxes, fontsize=fs)\n",
    "    ax.text(0.99, 0.85, f\"target={target}\", ha=\"right\", va=\"top\", transform=ax.transAxes, fontsize=fs)\n",
    "    ax.text(0.99, 0.8, f\"dev={dev}\", ha=\"right\", va=\"top\", transform=ax.transAxes, fontsize=fs)\n",
    "    # break\n",
    "\n",
    "    cc += 1\n",
    "    # break\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5694248b-8adc-4213-86f0-a530cf3d0264",
   "metadata": {},
   "outputs": [],
   "source": [
    "closest_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4635758-fdde-4a93-8534-ab4cc76f9b23",
   "metadata": {},
   "outputs": [],
   "source": [
    "closest_smiles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dcb79e6-77d9-4f59-9116-d60e7561eab5",
   "metadata": {},
   "source": [
    "### Experiment pca"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dfb0913-0c1d-4097-b787-5e9aebae2b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_w = pca.transform(X_exp)\n",
    "kept_feff_spectra_w = pca.transform(kept_feff_spectra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6ffc219-fc95-4276-9938-bfc5e25d0289",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(2, 2), sharey=True)\n",
    "\n",
    "remove_axis_spines(ax, visible=True)\n",
    "\n",
    "density_scatter(w[:, 0], w[:, 1], ax, bins=20, s=0.2, cmap=\"Blues\", rasterized=True)\n",
    "\n",
    "istar = data[\"functional_groups\"].index(\"Amine\")\n",
    "condition = (data[\"Y_test\"][:, istar] == 1)\n",
    "where = np.where(condition)[0]\n",
    "density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "\n",
    "for ii, point in enumerate(exp_w):\n",
    "    pred = mu[ii, feff_FG_index].item()\n",
    "    print(pred)\n",
    "    if pred == 1:\n",
    "        ax.scatter(point[0], point[1], c=\"black\", marker=\"x\", s=5)\n",
    "    else:\n",
    "        ax.scatter(point[0], point[1], c=\"black\", marker=\"o\", s=5)\n",
    "\n",
    "# density_scatter(kept_feff_spectra_w[:, 0], kept_feff_spectra_w[:, 1], ax, s=0.3, cmap=\"magma\")\n",
    "\n",
    "# ax.text(0.05, 0.95, f\"({label})\", ha=\"left\", va=\"top\", transform=ax.transAxes)\n",
    "\n",
    "# plt.subplots_adjust(wspace=0.1)\n",
    "plt.show()\n",
    "# plt.savefig(\"figures/n_pca.svg\", bbox_inches=\"tight\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "600a8f47-fea4-4574-ac7e-600d8bb24512",
   "metadata": {},
   "source": [
    "# Oxygen exploration and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679b329a-0454-432c-bec0-3f700a076e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "ELEMENT = \"O\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb02015f-f540-4988-8c98-119ecdf03c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble = Ensemble.from_path(f\"data/23-12-06_torch_models/{ELEMENT}\")\n",
    "data = get_data(elements=ELEMENT)\n",
    "grid = np.loadtxt(f\"data/{ELEMENT.lower()}_grid.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47b835b4-1881-4c6d-b575-146e42f8a87b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_spectrum = data[\"X_train\"].mean(axis=0).squeeze()\n",
    "for ii, (e, s) in enumerate(zip(grid, mean_spectrum)):\n",
    "    if s > 1:\n",
    "        break\n",
    "print(grid[ii])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7219109-3b92-455a-b5c3-4ea42bc39298",
   "metadata": {},
   "source": [
    "## PCA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1d70a8a-7081-44d6-8077-6d0ec301be72",
   "metadata": {},
   "source": [
    "See `Nitrogen-Oxygen preliminary analysis figures` section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bc7873f-bb9f-4f76-8671-b161f8ca636d",
   "metadata": {},
   "source": [
    "# Nitrogen-Oxygen preliminary analysis figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f2c9505-af90-41ed-9ebd-91f2f34090c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 4\n",
    "\n",
    "ELEMENT = \"N\"\n",
    "data_N = get_data(elements=ELEMENT)\n",
    "grid_N = np.loadtxt(f\"data/{ELEMENT.lower()}_grid.txt\")\n",
    "interesting_functional_groups_N = [\n",
    "    # None,\n",
    "    \"Amine\",\n",
    "    \"Amide\",\n",
    "    \"Nitrile\",\n",
    "    \"Heteroaromatic\",\n",
    "    # \"Secondary_aliph_amine\",\n",
    "    # \"Tertiary_aliph_amine\",\n",
    "]\n",
    "pca = PCA(2)\n",
    "w_N = pca.fit_transform(data_N[\"X_test\"])\n",
    "\n",
    "smiles = [\"N(C)(C)(C)\", \"CC(=O)N(C)(C)\", \"CC#N\", \"c1ccccn1\"]\n",
    "svgs = Chem.Draw.MolsToGridImage([Chem.MolFromSmiles(s) for s in smiles], useSVG=True)\n",
    "with open(\"figures/fig_pca_3-nitrogen-oxygen/n_motifs.svg\", \"w\") as f:\n",
    "    f.write(svgs.data)\n",
    "\n",
    "\n",
    "ELEMENT = \"O\"\n",
    "data_O = get_data(elements=ELEMENT)\n",
    "grid_O = np.loadtxt(f\"data/{ELEMENT.lower()}_grid.txt\")\n",
    "interesting_functional_groups_O = [\n",
    "    # None,\n",
    "    \"Alcohol\",\n",
    "    \"Epoxide\",\n",
    "    \"Carboxylic_acid_derivative\",\n",
    "    \"Ketone\",\n",
    "]\n",
    "pca = PCA(2)\n",
    "w_O = pca.fit_transform(data_O[\"X_test\"])\n",
    "\n",
    "smiles = [\"CO\",  \"C1CO1\", \"CC(=O)O\", \"CC(=O)C\"]\n",
    "svgs = Chem.Draw.MolsToGridImage([Chem.MolFromSmiles(s) for s in smiles], useSVG=True)\n",
    "with open(\"figures/fig_pca_3-nitrogen-oxygen/o_motifs.svg\", \"w\") as f:\n",
    "    f.write(svgs.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d6aedbc-9fed-4a6f-9f86-5d2dc1432b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = [\"i\", \"ii\", \"iii\", \"iv\", \"v\", \"vi\", \"vii\", \"viii\"]\n",
    "fig, axs_all = plt.subplots(2, N, figsize=(2 * N, 4), sharey=True, sharex=True)\n",
    "\n",
    "cc = 0\n",
    "\n",
    "axs = axs_all[0]\n",
    "w = w_N\n",
    "data = data_N\n",
    "ifg = interesting_functional_groups_N\n",
    "for fg, ax in zip(ifg, axs):\n",
    "    remove_axis_spines(ax, visible=True)\n",
    "    density_scatter(w[:, 0], w[:, 1], ax, bins=20, s=0.2, cmap=\"Blues\", rasterized=True)\n",
    "    if fg is not None:\n",
    "        index = data[\"functional_groups\"].index(fg)\n",
    "        condition = (data[\"Y_test\"][:, index] == 1)\n",
    "        # Plot where one nitrogen\n",
    "        condition_only_one = [Counter(xx.lower())[\"n\"] == 1 for xx in data[\"smiles_test\"]]\n",
    "        where = np.where(condition & condition_only_one)[0]\n",
    "        density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "        ax.text(0.05, 0.95, f\"({labels[cc]})\", ha=\"left\", va=\"top\", transform=ax.transAxes)\n",
    "        cc += 1\n",
    "\n",
    "axs = axs_all[1]\n",
    "w = w_O\n",
    "data = data_O\n",
    "ifg = interesting_functional_groups_O\n",
    "for fg, ax in zip(ifg, axs):\n",
    "    remove_axis_spines(ax, visible=True)\n",
    "    density_scatter(w[:, 0], w[:, 1], ax, bins=20, s=0.2, cmap=\"Reds\", rasterized=True)\n",
    "    if fg is not None:\n",
    "        index = data[\"functional_groups\"].index(fg)\n",
    "        condition = (data[\"Y_test\"][:, index] == 1)\n",
    "        # Plot where one oxygen\n",
    "        if fg == \"Carboxylic_acid_derivative\":\n",
    "            condition_only_one = [Counter(xx.lower())[\"o\"] == 2 for xx in data[\"smiles_test\"]]\n",
    "        else:\n",
    "            condition_only_one = [Counter(xx.lower())[\"o\"] == 1 for xx in data[\"smiles_test\"]]\n",
    "        where = np.where(condition & condition_only_one)[0]\n",
    "        density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "        ax.text(0.05, 0.95, f\"({labels[cc]})\", ha=\"left\", va=\"top\", transform=ax.transAxes)\n",
    "        cc += 1\n",
    "\n",
    "axs[0].set_xlabel(r\"$z_1$~[a.u.]\")\n",
    "axs[0].set_ylabel(r\"$z_2$~[a.u.]\")\n",
    "\n",
    "plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
    "plt.show()\n",
    "# plt.savefig(\"figures/fig_pca_3-nitrogen-oxygen/pca.svg\", bbox_inches=\"tight\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cb69110-2523-4f0f-9450-5bb4817d789e",
   "metadata": {},
   "source": [
    "## Spectra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0c780e4-78f6-42c4-8630-e39968a7056c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mult = 4\n",
    "\n",
    "fig, axs = plt.subplots(2, 1, figsize=(1.5, 4), sharey=True)\n",
    "\n",
    "# remove_axis_spines(ax, visible=True)\n",
    "\n",
    "cmap = mpl.colormaps[\"rainbow\"].resampled(len(interesting_functional_groups))\n",
    "\n",
    "ax = axs[0]\n",
    "data = data_N\n",
    "ifg = interesting_functional_groups_N\n",
    "X = data_N[\"X_test\"]\n",
    "grid = np.loadtxt(\"data/N_grid.txt\")\n",
    "\n",
    "for ii, (fg,) in enumerate(zip(ifg)):\n",
    "    \n",
    "    index = data[\"functional_groups\"].index(fg)\n",
    "    condition = (data[\"Y_test\"][:, index] == 1)\n",
    "    condition_only_one = [Counter(xx.lower())[\"n\"] == 1 for xx in data[\"smiles_test\"]]\n",
    "    \n",
    "    where = np.where(condition & condition_only_one)[0]\n",
    "    \n",
    "    X_where = X[where, :]\n",
    "    mu = X_where.mean(axis=0)\n",
    "    sd = X_where.std(axis=0)\n",
    "    ax.plot(grid, mu - ii*mult, color=cmap(ii), zorder=-ii)\n",
    "    for jj in range(100, 110):\n",
    "        ax.plot(grid, X_where[jj, :] - ii * mult, color=cmap(ii), alpha = 0.5, linewidth=0.5)\n",
    "    # ax.fill_between(grid, mu + ii*mult - sd, mu + ii*mult + sd, linewidth=0, color=cmap(ii), alpha=0.2)\n",
    "\n",
    "ax.set_yticks([])\n",
    "set_grids(ax, grid=False)\n",
    "ax.set_ylabel(r\"$\\mu(E)$~[a.u.]\")\n",
    "# ax.set_xlabel(r\"$E$~[eV]\")\n",
    "ax.set_xticks([400, 410, 420, 430])\n",
    "ax.set_xlim([397, 431])\n",
    "\n",
    "ax = axs[1]\n",
    "data = data_O\n",
    "ifg = interesting_functional_groups_O\n",
    "X = data_O[\"X_test\"]\n",
    "grid = np.loadtxt(\"data/O_grid.txt\")\n",
    "\n",
    "for ii, (fg,) in enumerate(zip(ifg)):\n",
    "    \n",
    "    index = data[\"functional_groups\"].index(fg)\n",
    "    condition = (data[\"Y_test\"][:, index] == 1)\n",
    "    condition_only_one = [Counter(xx.lower())[\"o\"] == 1 for xx in data[\"smiles_test\"]]\n",
    "    \n",
    "    where = np.where(condition & condition_only_one)[0]\n",
    "    \n",
    "    X_where = X[where, :]\n",
    "    mu = X_where.mean(axis=0)\n",
    "    sd = X_where.std(axis=0)\n",
    "    ax.plot(grid, mu - ii*mult, color=cmap(ii), zorder=-ii)\n",
    "    for jj in range(100, 110):\n",
    "        ax.plot(grid, X_where[jj, :] - ii * mult, color=cmap(ii), alpha = 0.5, linewidth=0.5)\n",
    "    # ax.fill_between(grid, mu + ii*mult - sd, mu + ii*mult + sd, linewidth=0, color=cmap(ii), alpha=0.2)\n",
    "\n",
    "ax.set_yticks([])\n",
    "set_grids(ax, grid=False)\n",
    "# ax.set_ylabel(r\"$\\mu(E)$~[a.u.]\")\n",
    "ax.set_xlabel(r\"$E$~[eV]\")\n",
    "ax.set_xticks([530, 540, 550, 560])\n",
    "ax.set_xlim([529, 561])\n",
    "\n",
    "plt.subplots_adjust(wspace=0.1, hspace=0.15)\n",
    "plt.savefig(\"figures/fig_pca_3-nitrogen-oxygen/spectra.svg\", bbox_inches='tight', dpi=300)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b1688a1-f5e2-4096-ad0d-bb3228ee0bd6",
   "metadata": {},
   "source": [
    "## Analyze specific regions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2ab414f-84e6-42ba-a9d1-c708a5c4035b",
   "metadata": {},
   "source": [
    "### Secondary/Tertiary amines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0daf9ab-f0b4-468a-a2c5-1fe72b520a67",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_N[\"X_test\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702232a5-520c-4fa8-a889-36e5c5f4be9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(w_N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74ed5ad5-aab0-4982-a90e-64cde3c797ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(2, 2))\n",
    "\n",
    "w = w_O\n",
    "data = data_O\n",
    "ifg = interesting_functional_groups_O\n",
    "fg = \"Carboxylic_acid_derivative\"\n",
    "\n",
    "density_scatter(w[:, 0], w[:, 1], ax, bins=20, s=0.2, cmap=\"Reds\", rasterized=True)\n",
    "\n",
    "index = data[\"functional_groups\"].index(fg)\n",
    "condition = (data[\"Y_test\"][:, index] == 1)\n",
    "condition_only_one = [Counter(xx.lower())[\"o\"] == 2 for xx in data[\"smiles_test\"]]\n",
    "\n",
    "# condition_location = (w[:, 0] > 0) & (w[:, 0] < 4.5)\n",
    "condition_location = w[:, 0] > 4.5\n",
    "\n",
    "where = np.where(condition & condition_only_one & condition_location)[0]\n",
    "\n",
    "density_scatter(w[where, 0], w[where, 1], ax, s=0.3, cmap=\"viridis\", rasterized=True)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1fbdc0-8634-4f54-b09c-78a36e049a0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "condition = (data[\"Y_test\"][:, index] == 1)\n",
    "condition_only_one = [Counter(xx.lower())[\"o\"] == 2 for xx in data[\"smiles_test\"]]\n",
    "condition_location = w[:, 0] > 4.5\n",
    "where = np.where(condition & condition_only_one & condition_location)[0]\n",
    "smiles = [data[\"smiles_test\"][ii] for ii in where][:10]\n",
    "svgs = Chem.Draw.MolsToGridImage([Chem.MolFromSmiles(s) for s in smiles], useSVG=True)\n",
    "svgs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb8c7d26-b8fc-46cf-9696-1ac78d7e9ba3",
   "metadata": {},
   "source": [
    "# Model results and UQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a8e8642-8193-4d6c-940e-9c42a82ce529",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_containing_functional_groups = [\n",
    "    '1,2-Aminoalcohol',\n",
    "    'Lactam',\n",
    "    'Amide',\n",
    "    'Imidolactone',\n",
    "    'Heterocyclic',\n",
    "    'Hetero_N_basic_H',\n",
    "    'Amine',\n",
    "    'Hetero_N_nonbasic',\n",
    "    'Primary_arom_amine',\n",
    "    'Tertiary_aliph_amine',\n",
    "    'NH_aziridine',\n",
    "    'Heteroaromatic',\n",
    "    'Secondary_aliph_amine',\n",
    "    'Nitrile',\n",
    "]\n",
    "\n",
    "O_containing_functional_groups = [\n",
    "    '1,2-Aminoalcohol',\n",
    "    'Secondary_alcohol',\n",
    "    'Carbonic_acid_derivatives',\n",
    "    'Tertiary_alcohol',\n",
    "    'Lactam',\n",
    "    'Primary_alcohol',\n",
    "    'Aldehyde',\n",
    "    'Ketone',\n",
    "    'Carboxylic_acid_derivative',\n",
    "    'Epoxide',\n",
    "    'Imidolactone',\n",
    "    'Heterocyclic',\n",
    "    'Dialkylether',\n",
    "    'Phenol',\n",
    "    'Heteroaromatic',\n",
    "    'Hetero_O',\n",
    "    'Alcohol',\n",
    "]\n",
    "\n",
    "def get_statistics(\n",
    "    deviating_estimators,\n",
    "    functional_groups,\n",
    "    truth,\n",
    "    mu,\n",
    "    element,\n",
    "):\n",
    "    \n",
    "    cba_mean = []\n",
    "    cba_var = []\n",
    "    cba_min = []\n",
    "    cba_max = []\n",
    "    cba_counts = []\n",
    "    \n",
    "    special = {\"1,2-Aminoalcohol\": []}\n",
    "    \n",
    "    for ii, n_dev in enumerate(range(0, 11)):\n",
    "        tmp_cba = []\n",
    "        tmp_weights = []\n",
    "        where = np.where(deviating_estimators.flatten() == n_dev)[0]\n",
    "        cba_counts.append(len(where))\n",
    "        \n",
    "        for jj, fg in enumerate(functional_groups):\n",
    "            if element == \"O\" and fg not in O_containing_functional_groups:\n",
    "                continue\n",
    "            if element == \"N\" and fg not in N_containing_functional_groups:\n",
    "                continue\n",
    "\n",
    "            where = np.where(deviating_estimators[:, jj] == n_dev)[0]\n",
    "            if len(where) < 10:\n",
    "                continue\n",
    "\n",
    "            acc = balanced_accuracy_score(truth[where, jj].flatten(), mu[where, jj].flatten())\n",
    "\n",
    "            if element == \"N\" and fg == \"1,2-Aminoalcohol\":\n",
    "                special[\"1,2-Aminoalcohol\"].append(acc)\n",
    "                continue\n",
    "\n",
    "            tmp_weights.append(len(where))\n",
    "            tmp_cba.append(acc)\n",
    "\n",
    "        mean = np.average(tmp_cba, weights=tmp_weights)\n",
    "        cba_mean.append(mean)\n",
    "        var = np.average((mean - np.array(tmp_cba))**2, weights=tmp_weights)\n",
    "        cba_var.append(var)\n",
    "        cba_min.append(np.min(tmp_cba))\n",
    "        cba_max.append(np.max(tmp_cba))\n",
    "\n",
    "    return {\n",
    "        \"mean\": np.array(cba_mean),\n",
    "        \"var\": np.array(cba_var),\n",
    "        \"min\": np.array(cba_min),\n",
    "        \"max\": np.array(cba_max),\n",
    "        \"counts\": np.array(cba_counts),\n",
    "        **special\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04892a4b-b70d-49ac-883d-6b2f108d1fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data = get_data(elements=\"O\")\n",
    "# ensemble = Ensemble.from_path(\"data/23-12-06_torch_models/O\")\n",
    "# preds = ensemble.predict(data[\"X_test\"])\n",
    "# mu = preds.round().mean(axis=0).round()\n",
    "# deviating_estimators = get_deviating_estimators(preds)\n",
    "# truth = data[\"Y_test\"]\n",
    "# functional_groups = data[\"functional_groups\"]\n",
    "# jj = functional_groups.index(\"1,2-Aminoalcohol\")\n",
    "# balanced_accuracy_score(truth[:, jj], mu[:, jj])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02c974c0-1010-4cd1-aed5-4f1cecfce9b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = {}\n",
    "ELEMENTS = [\"C\", \"N\", \"O\", \"CNO\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62bf8860-18a5-45d0-a948-a0090cc4b8f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use this one for main results\n",
    "target = \"data/23-04-26-ml-data\"\n",
    "estimator_directory = Path(\"data/23-12-06_torch_models\")\n",
    "\n",
    "# use this one for the cutoff-8 results\n",
    "# target = \"data/23-05-11-ml-data-CUTOFF8\"\n",
    "# estimator_directory = Path(\"data/23-12-06_torch_models/cutoff8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "940b8cea-6c35-4545-ae23-7201324e36d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for elements in tqdm.tqdm(ELEMENTS):\n",
    "\n",
    "    data = get_data(target=target, elements=elements)\n",
    "    p = \"-\".join([e for e in elements])\n",
    "    ensemble = Ensemble.from_path(estimator_directory / p)\n",
    "\n",
    "    preds = ensemble.predict(data[\"X_test\"])\n",
    "    mu = preds.round().mean(axis=0).round()\n",
    "    deviating_estimators = get_deviating_estimators(preds)\n",
    "    truth = data[\"Y_test\"]\n",
    "    functional_groups = data[\"functional_groups\"]\n",
    "    print(elements, functional_groups)\n",
    "\n",
    "    stats[elements] = get_statistics(\n",
    "        deviating_estimators,\n",
    "        functional_groups,\n",
    "        truth,\n",
    "        mu,\n",
    "        elements,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a273ab-19b3-45d9-a9c4-70758b071ed7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_kwargs = {\n",
    "    'linewidth': 1.0,\n",
    "    'marker': 's',\n",
    "    'ms': 2.0,\n",
    "    'capthick': 0.3,\n",
    "    'capsize': 2.0,\n",
    "    'elinewidth': 0.3\n",
    "}\n",
    "COLORS = [\"black\", \"blue\", \"red\", \"grey\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d810458-6c4f-4e20-a7d2-c4bf825952e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 4, figsize=(6, 2), sharey=True)\n",
    "\n",
    "for ii, (ax, color, (key, value)) in enumerate(zip(axs, COLORS, stats.items())):\n",
    "\n",
    "    cba_mean = value[\"mean\"]\n",
    "    cba_var = value[\"var\"]\n",
    "    cba_std = np.sqrt(cba_var)\n",
    "    cba_min = value[\"min\"]\n",
    "    cba_max = value[\"max\"]\n",
    "    cba_counts = value[\"counts\"]\n",
    "    \n",
    "    ax.errorbar(\n",
    "        range(0, 11),\n",
    "        cba_mean,\n",
    "        yerr=[cba_mean-cba_min, cba_max-cba_mean],\n",
    "        color=color,\n",
    "        label=key,\n",
    "        **plot_kwargs\n",
    "    )\n",
    "\n",
    "    if ii == 1:  # N\n",
    "        ax.errorbar(\n",
    "            range(0, 11),\n",
    "            value[\"1,2-Aminoalcohol\"],\n",
    "            color=\"purple\",\n",
    "            alpha=1,\n",
    "            label=\"AA\",\n",
    "            **plot_kwargs\n",
    "        )\n",
    "\n",
    "    ax.legend(fancybox=True, fontsize=10)\n",
    "\n",
    "    if ii > 0:\n",
    "        ax2.set_yticklabels([])\n",
    "    \n",
    "    ax2 = ax.twinx()\n",
    "    ax2.bar(range(0, 11), np.log10(cba_counts), color=color, alpha=0.2)\n",
    "    ax2.set_ylim(bottom=3.4, top=7.1)\n",
    "    ax2.set_yticks([4, 5, 6, 7])\n",
    "    ax.tick_params(which=\"both\", direction=\"in\")\n",
    "    ax2.tick_params(which=\"both\", direction=\"in\")\n",
    "\n",
    "    if ii == 3:\n",
    "        ax2.set_ylabel(r\"$\\log_{10}$ Counts\")\n",
    "    if ii == 0:\n",
    "        ax.set_ylabel(r\"CBA\")\n",
    "\n",
    "ax.set_xlabel(r\"$N_\\mathrm{d}$\")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# plt.savefig(\"figures/fig_nd/fig_nd_2.svg\", dpi=300, bbox_inches=\"tight\")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45510a3c-c556-473f-858f-cf04b3bfba71",
   "metadata": {},
   "source": [
    "# Multimodal advantage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "933e4c96-a0d3-47ae-8b76-905084897fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use this one for main results\n",
    "# target = \"data/23-04-26-ml-data\"\n",
    "# estimator_directory = Path(\"data/23-12-06_torch_models/solo_estimators\")\n",
    "\n",
    "# use this one for the cutoff-8 results\n",
    "target = \"data/23-05-11-ml-data-CUTOFF8\"\n",
    "estimator_directory = Path(\"data/23-12-06_torch_models/solo_estimators_cutoff8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f4f57e-a108-4f53-b4c7-7041315665fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_CNO = get_data(target=target, elements=\"CNO\")\n",
    "data_C = get_data(target=target, elements=\"C\")\n",
    "data_N = get_data(target=target, elements=\"N\")\n",
    "data_O = get_data(target=target, elements=\"O\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbe68dd8-108e-4c3a-8afc-99109afdeb49",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_map(input_data, element):\n",
    "    if element == \"C\":\n",
    "        return input_data[:, :200]\n",
    "    elif element == \"N\":\n",
    "        return input_data[:, 200:400]\n",
    "    elif element == \"O\":\n",
    "        return input_data[:, 400:]\n",
    "    elif element == \"CNO\":\n",
    "        return input_data\n",
    "    raise ValueError(f\"invalid element {element}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9893ca-84cc-4221-b62f-e910f2c67efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator_CNO_C = Estimator.from_path(estimator_directory / \"C-N-O_only_C\")\n",
    "estimator_CNO_N = Estimator.from_path(estimator_directory / \"C-N-O_only_N\")\n",
    "estimator_CNO_O = Estimator.from_path(estimator_directory / \"C-N-O_only_O\")\n",
    "estimator_CNO = Estimator.from_path(estimator_directory / \"C-N-O\")\n",
    "estimator_C = Estimator.from_path(estimator_directory / \"C\")\n",
    "estimator_N = Estimator.from_path(estimator_directory / \"N\")\n",
    "estimator_O = Estimator.from_path(estimator_directory / \"O\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "854d08d3-8d4d-4b94-87d9-4020ca76a621",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_CNO_C = estimator_CNO_C.predict(index_map(data_CNO[\"X_test\"], element=\"C\"))\n",
    "pred_CNO_N = estimator_CNO_N.predict(index_map(data_CNO[\"X_test\"], element=\"N\"))\n",
    "pred_CNO_O = estimator_CNO_O.predict(index_map(data_CNO[\"X_test\"], element=\"O\"))\n",
    "pred_CNO = estimator_CNO.predict(index_map(data_CNO[\"X_test\"], element=\"CNO\"))\n",
    "\n",
    "pred_C = estimator_C.predict(data_C[\"X_test\"])\n",
    "pred_N = estimator_N.predict(data_N[\"X_test\"])\n",
    "pred_O = estimator_O.predict(data_O[\"X_test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a8e96dc-dea9-4b41-858e-e9311177eb98",
   "metadata": {},
   "source": [
    "First, it's useful to have the overall CBA score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d7127d4-ce30-4b95-95ae-bbad77c19244",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_test = data_CNO[\"Y_test\"]\n",
    "balanced_accuracy_score(Y_test.flatten(), pred_CNO.round().flatten())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aeb41d3-f078-4e7c-a6bd-a49a2ec42314",
   "metadata": {},
   "source": [
    "Then we continue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a695a141-8a46-4247-999b-99247517d29b",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_containing_functional_groups = [\n",
    "    '1,2-Aminoalcohol',\n",
    "    'Lactam',\n",
    "    'Amide',\n",
    "    'Imidolactone',\n",
    "    'Heterocyclic',\n",
    "    'Hetero_N_basic_H',\n",
    "    'Amine',\n",
    "    'Hetero_N_nonbasic',\n",
    "    'Primary_arom_amine',\n",
    "    'Tertiary_aliph_amine',\n",
    "    'NH_aziridine',\n",
    "    'Heteroaromatic',\n",
    "    'Secondary_aliph_amine',\n",
    "    'Nitrile',\n",
    "]\n",
    "\n",
    "O_containing_functional_groups = [\n",
    "    '1,2-Aminoalcohol',\n",
    "    'Secondary_alcohol',\n",
    "    'Carbonic_acid_derivatives',\n",
    "    'Tertiary_alcohol',\n",
    "    'Lactam',\n",
    "    'Primary_alcohol',\n",
    "    'Aldehyde',\n",
    "    'Ketone',\n",
    "    'Carboxylic_acid_derivative',\n",
    "    'Epoxide',\n",
    "    'Imidolactone',\n",
    "    'Heterocyclic',\n",
    "    'Dialkylether',\n",
    "    'Phenol',\n",
    "    'Heteroaromatic',\n",
    "    'Hetero_O',\n",
    "    'Alcohol',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "446334f2-1228-4867-934a-e884df2947a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {\n",
    "    \"fg\": data_CNO[\"functional_groups\"], \"C\": [], \"N\": [], \"O\": [], \"CNO\": [] # \"CNO_C\": [], \"CNO_N\": [], \"CNO_O\": [],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a0b0e8-baef-4a0e-b93b-542517f5790b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ii, fg in enumerate(data_CNO[\"functional_groups\"]):\n",
    "    Y_test = data_CNO[\"Y_test\"]\n",
    "    a_CNO = balanced_accuracy_score(Y_test[:, ii], pred_CNO.round()[:, ii])\n",
    "    a_C = balanced_accuracy_score(Y_test[:, ii], pred_CNO_C.round()[:, ii])\n",
    "    a_N = balanced_accuracy_score(Y_test[:, ii], pred_CNO_N.round()[:, ii])\n",
    "    a_O = balanced_accuracy_score(Y_test[:, ii], pred_CNO_O.round()[:, ii])\n",
    "    \n",
    "    results[\"C\"].append(a_C)\n",
    "    results[\"N\"].append(a_N)\n",
    "    results[\"O\"].append(a_O)\n",
    "    results[\"CNO\"].append(a_CNO)\n",
    "    # print(f\"{fg:<30} {a_C:.03f} {a_N:.03f} {a_O:.03f} {a_CNO:.03f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26dccc66-7a34-43d2-9d6f-9551ec616888",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(results)\n",
    "df[\"Best SM\"] = df[[\"C\", \"N\", \"O\"]].max(axis=1)\n",
    "df[\"d\"] = df[\"CNO\"] - df[\"Best SM\"]\n",
    "df = df.sort_values(\"d\", ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ff74fb-961c-4fc9-89f2-c271f808c682",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_plot(df):\n",
    "    \n",
    "    # fig, axs = plt.subplots(2, 1, figsize=(8, 2), gridspec_kw={'height_ratios': [1, 3]}, sharex=True)\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(7, 2))\n",
    "\n",
    "    x = [ii for ii in range(len(df.index))]\n",
    "\n",
    "    ax.scatter(x, df[\"CNO\"], color=\"grey\", label=r\"$\\mathcal{D}_\\mathrm{CNO}$\")\n",
    "\n",
    "\n",
    "    s = 10\n",
    "    ax.scatter(x, df[\"C\"], color=\"black\", s=s, label=r\"$\\mathcal{D}_\\mathrm{CNO}^\\mathrm{[C]}$\", linewidth=0.5)\n",
    "\n",
    "    for ii, (fg, xx) in enumerate(zip(df[\"fg\"], x)):\n",
    "        ax.scatter(\n",
    "            xx, df[df[\"fg\"] == fg][\"N\"],\n",
    "            s=s, label=r\"$\\mathcal{D}_\\mathrm{CNO}^\\mathrm{[N]}$\" if ii == 0 else None, facecolors=\"white\" if fg not in N_containing_functional_groups else \"blue\", edgecolors=\"blue\",\n",
    "            linewidth=.5\n",
    "        )\n",
    "\n",
    "    for ii, (fg, xx) in enumerate(zip(df[\"fg\"], x)):\n",
    "        ax.scatter(\n",
    "            xx, df[df[\"fg\"] == fg][\"O\"],\n",
    "            s=s, label=r\"$\\mathcal{D}_\\mathrm{CNO}^\\mathrm{[O]}$\" if ii == 0 else None, facecolors=\"white\" if fg not in O_containing_functional_groups else \"red\", edgecolors=\"red\",\n",
    "            linewidth=.5\n",
    "        )\n",
    "\n",
    "\n",
    "    labels = [xx.replace(\"_\", \" \") for xx in df[\"fg\"]]\n",
    "    ax.set_xticks(x)\n",
    "    ax.set_xticklabels(labels, rotation=45, ha='right', fontsize=8)\n",
    "\n",
    "    ax.tick_params(which='both', direction='in', bottom=True, left=True, top=False, right=False)\n",
    "    for xx in x:\n",
    "        ax.axvline(xx, zorder=-1, linewidth=0.5, color=\"black\", alpha=0.2)\n",
    "\n",
    "    ax.set_ylabel(\"CBA\", fontsize=10)\n",
    "\n",
    "    # ax.legend(frameon=True, ncol=4, fontsize=6, loc=\"lower center\", bbox_to_anchor=(1, -1))\n",
    "    \n",
    "    \n",
    "\n",
    "    \n",
    "    ax2 = ax.twinx() \n",
    "    \n",
    "    ax2.scatter(x, df[\"d\"] * 100, zorder=-1, marker=\"^\", color=\"green\", s=5)\n",
    "    # ax.axhline(df[\"d\"].mean() * 100, color=\"grey\", zorder=-1, linewidth=0.5)\n",
    "    # ax.bar(x, df[\"d8\"] * 100, color=\"purple\", width=0.5, alpha=1)\n",
    "    # ax.axhline(df[\"d8\"].mean() * 100, color=\"purple\", zorder=-1, linewidth=0.5)\n",
    "    ax2.set_ylabel(\"Adv (\\%)\", fontsize=10, color=\"green\")\n",
    "    ax2.tick_params(which='both', direction='in', bottom=False, left=False, top=True, right=True)\n",
    "    \n",
    "    return ax, ax2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57572c2f-c60b-458c-bdc4-48d28495d1ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "axs = make_plot(df)\n",
    "\n",
    "axs[0].set_yticks([0.5, 0.6, 0.7, 0.8, 0.9, 1.0])\n",
    "yticks = axs[0].get_yticks()\n",
    "axs[0].set_ylim(yticks[0] - 0.02, yticks[-1] + 0.02)\n",
    "axs[1].set_yticks([0, 3, 6])\n",
    "\n",
    "# axs[0].set_yticks([0.7, 0.8, 0.9, 1.0])\n",
    "# axs[0].set_ylim(0.68, 1.02)\n",
    "# axs[1].set_yticks([0, 3, 6])\n",
    "\n",
    "plt.savefig(\"figures/fig_multimodal_advantage/multimodal_advantage-cutoff8.svg\", bbox_inches=\"tight\", dpi=300)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56cb7fff-062e-42d0-b418-33cad9602163",
   "metadata": {},
   "outputs": [],
   "source": [
    "-2 % 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a062d0-3d9c-4e97-b1e1-d8d3904fadb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddd = [1, 2, 3]\n",
    "ddd.insert(0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "959d45ce-97cb-4342-983e-ead36bd28e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79b480e9-0c0f-4c5d-aa35-5503cfc7f604",
   "metadata": {},
   "outputs": [],
   "source": [
    "del ddd[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f7b4ba-6acb-426c-bcfc-c89821db6ecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6b33517-2151-42d5-86f4-b8d5a5590d96",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}