1. Data preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a4278ec6",
   "metadata": {},
   "source": [
    "# Data preprocessing\n",
    "\n",
    "The data about the molecules is stored in the database at path `data/trian.db`. It can be accessed through the `ase` library. Each molecule is stored as a row which contains:\n",
    "\n",
    "- `row.natoms` – number of atoms in the molecule\n",
    "- `row.symbols` – `np.array` of atom names\n",
    "- `row.positions` – `np.array` of atom positions in 3d space\n",
    "- `row.data.energy` – a list with the energy of the molecule\n",
    "\n",
    "The final data for the neural network will be a table where each row correcponds to a molecule, and each column is the number of each type of bond. To get from a position matrix to a bond type vector the following steps have to be made:\n",
    "- function to convert a position matrix to distance matrix\n",
    "- function to convert from the dataset units to real life units (picometers)\n",
    "- function to convert a distance matrix to bond type matrix\n",
    "- function to count each bond type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fc7d2239",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML\n",
    "\n",
    "from ase.db import connect\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from sklearn.cluster import KMeans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0000d40f",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_db = connect(\"./data/train.db\")\n",
    "\n",
    "TOTAL_MOLECULES = 138365"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb2dd159",
   "metadata": {},
   "source": [
    "## Position matrix to distance matrix\n",
    "\n",
    "This can be done with multiple for loops, but leveraging the linear algebra knowledge would increase the speed of the code and its simplicity. In this section a feature of numpy will be used called broadcasting. It is when the matricies are not the exact same size however they can be broad casted such that they have the same shape.\n",
    "\n",
    "In this problems we have 3 dimensions $(x, y, z)$, however for understanding purposes let's imagine that we have only $x$. Let $X$ be a vector of $x$-positions of each molecule\n",
    "\n",
    "$$ X = \n",
    "\\begin{bmatrix}\n",
    "x_1 \\\\\n",
    "x_2 \\\\\n",
    "\\vdots \\\\\n",
    "x_n\n",
    "\\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "Now, let's subtract $X^T$ from $X$. Normally we wouldn't be able to do that, however with broadcasting we can, which will yeald us\n",
    "\n",
    "$$\n",
    "X - X^T = \n",
    "\\begin{bmatrix}\n",
    "x_1    & x_1    & \\dots & x_1 \\\\\n",
    "x_2    & x_2    & \\dots & x_2 \\\\\n",
    "\\vdots & \\vdots & \\dots & \\vdots \\\\\n",
    "x_n    & x_n    & \\dots & x_n\n",
    "\\end{bmatrix}\n",
    "-\n",
    "\\begin{bmatrix}\n",
    "x_1    & x_2    & \\dots & x_n \\\\\n",
    "x_1    & x_2    & \\dots & x_n \\\\\n",
    "\\vdots & \\vdots & \\dots & \\vdots \\\\\n",
    "x_1    & x_2    & \\dots & x_n\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "x_1 - x_1 & x_1 - x_2 & \\dots & x_1 - x_n \\\\\n",
    "x_2 - x_1 & x_2 - x_2 & \\dots & x_2 - x_n \\\\\n",
    "\\vdots    & \\vdots    & \\dots & \\vdots    \\\\\n",
    "x_n - x_1 & x_n - x_2 & \\dots & x_n - x_n \n",
    "\\end{bmatrix},\n",
    "$$\n",
    "\n",
    "this is the distance matrix for the $x$-axis. \n",
    "\n",
    "Repeating the same for the $Y$ and $Z$ and then taking the normal will give the distance matrix.\n",
    "\n",
    "$$\n",
    "\\text{norm}(\n",
    "\\begin{bmatrix}\n",
    "X - X^T & Y - Y^T & Z - Z^T\n",
    "\\end{bmatrix}\n",
    "\\big) = \\\\\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "\\sqrt{(x_1 - x_1)^2+(y_1 - y_1)^2+(z_1 - z_1)^2} & \\sqrt{(x_1 - x_2)^2+(y_1 - y_2)^2+(z_1 - z_2)^2} & \\dots & \\sqrt{(x_1 - x_n)^2+(y_1 - y_n)^2+(z_1 - z_n)^2} \\\\\n",
    "\\sqrt{(x_2 - x_1)^2+(y_2 - y_1)^2+(z_2 - z_1)^2} & \\sqrt{(x_2 - x_2)^2+(y_2 - y_2)^2+(z_2 - z_2)^2} & \\dots & \\sqrt{(x_2 - x_n)^2+(y_2 - y_n)^2+(z_2 - z_n)^2} \\\\\n",
    "\\vdots    & \\vdots    & \\dots & \\vdots    \\\\\n",
    "\\sqrt{(x_n - x_1)^2+(y_n - y_1)^2+(z_n - z_1)^2} & \\sqrt{(x_n - x_2)^2+(y_n - y_2)^2+(z_n - z_2)^2} & \\dots & \\sqrt{(x_n - x_n)^2+(y_n - y_n)^2+(z_n - z_n)^2} \n",
    "\\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b393d34c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 3., 4.],\n",
       "       [3., 0., 5.],\n",
       "       [4., 5., 0.]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def position_to_distance_matrix(position_matrix):\n",
    "    # number of dimensitons (= 3)\n",
    "    dims = position_matrix.shape[1]\n",
    "\n",
    "    # expand the shape of the matrix so that they are broadcastable\n",
    "    m1 = position_matrix.reshape((1, -1, dims))\n",
    "    m2 = position_matrix.reshape((-1, 1, dims))\n",
    "\n",
    "    return np.linalg.norm(m1 - m2, axis=-1)\n",
    "    \n",
    "position_to_distance_matrix(np.array([\n",
    "    [0, 0, 0],\n",
    "    [0, 0, 3],\n",
    "    [0, 4, 0]\n",
    "]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3137c7b",
   "metadata": {},
   "source": [
    "Works good."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec2a186e",
   "metadata": {},
   "source": [
    "## Converting units to picometers\n",
    "\n",
    "The units in the dataset aren't specified. So, we will have to figure out how to convert to them emperically. We'll do that by looking at the distribution of lengths of `C-H` bonds. Since, they can have only one bond, the carbon hydrogen pairs can be divided into two groups one with a bond, and one without.\n",
    "\n",
    "Based on covalent radii, `C-H` bond has an average length of $R_C + R_H = 31 + 76 = 107$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "34380939",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 138365/138365 [00:45<00:00, 3063.93it/s]\n"
     ]
    }
   ],
   "source": [
    "CH_average_bond_length_pm = 31 + 76\n",
    "CH_bond_lengths = []\n",
    "\n",
    "for i in tqdm(range(1, TOTAL_MOLECULES+1)):\n",
    "    row = train_db.get(i)\n",
    "    symbols = row.symbols\n",
    "    positions = row.positions\n",
    "\n",
    "    distance_matrix = position_to_distance_matrix(positions)\n",
    "\n",
    "    for i, C in enumerate(symbols):\n",
    "        if C != 'C': continue\n",
    "        for j, H in enumerate(symbols):\n",
    "            if H != 'H': continue\n",
    "            CH_bond_lengths.append(distance_matrix[i][j])\n",
    "            \n",
    "CH_bond_lengths = np.array(CH_bond_lengths)\n",
    "CH_convertion_to_pm = CH_average_bond_length_pm / CH_bond_lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "122a72ba",
   "metadata": {},
   "source": [
    "Let's cluster the convertion rates using kmeans algorithm with 2 means. If atoms are close, then they are most likely bonded, the bond convertion is the inverse of that, hence we need the bigger number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e1cc2508",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "27.152678863047377"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kmeans = KMeans(n_clusters=2, random_state=0).fit(CH_convertion_to_pm.reshape(-1,1))\n",
    "\n",
    "mean_converion_rate = max(kmeans.cluster_centers_[:, 0])\n",
    "\n",
    "mean_converion_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8c0188ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXUAAAEDCAYAAADKhpQUAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAQK0lEQVR4nO3df4xlZX3H8fenuNCmkFLdqRJ+jVKsEatIt6jVEGJig2CkTWkLGlFDu4mRVlttBZug1TRBk9pWUclWKGIVf1Rrt7rW0kqCNhXZxQX2R9VVt2HJ1l1AQKKpbv32j3sWb4aZuXd2786588z7ldzM+fHMvd88y3x45jnnPJOqQpLUhp/quwBJ0uQY6pLUEENdkhpiqEtSQwx1SWqIoS5JDek11JNcn2Rfkm1jtv+dJDuSbE/ykSNdnyStNOnzPvUk5wCPADdW1TNGtD0d+Djwwqr6bpJfqKp9y1GnJK0UvY7Uq+pW4IHhY0lOS/IvSbYk+WKSp3Wnfh94b1V9t/teA12S5pjGOfUNwB9U1a8AbwTe1x1/KvDUJP+R5MtJzuutQkmaUo/ru4BhSY4Ffg34RJKDh4/pvj4OOB04FzgJuDXJL1fVg8tcpiRNrakKdQa/OTxYVWfOc24PcFtV/Qj4dpKvMwj525exPkmaalM1/VJVDzMI7N8GyMCzutOfZjBKJ8laBtMx3+qhTEmaWn3f0ngT8J/ALyXZk+Qy4OXAZUnuBLYDF3bNPw/cn2QHcAvwJ1V1fx91S9K06vWWRknSZE3V9Isk6fD0dqF07dq1NTs729fHS9KKtGXLlvuqamah872F+uzsLJs3b+7r4yVpRUry34udd/pFkhpiqEtSQwx1SWqIoS5JDTHUJakhhrokNcRQl6SGGOqS1BBDXZIaMm3rqS+L2Ss+++j27qsv6LESSZosR+qS1BBDXZIaYqhLUkNWzZz68Dy6JLXKkbokNcRQl6SGGOqS1BBDXZIaYqhLUkMMdUlqiKEuSQ0x1CWpIYa6JDVkZKgnOTnJLUl2JNme5HXztDk3yUNJtnavq45MuZKkxYyzTMAB4A1VdUeS44AtSW6uqh1z2n2xql4y+RIlSeMaOVKvqr1VdUe3/T1gJ3DikS5MkrR0S5pTTzILPBu4bZ7Tz0tyZ5LPJTljge9fn2Rzks379+9ferWSpEWNHepJjgU+Cby+qh6ec/oO4NSqehbwHuDT871HVW2oqnVVtW5mZuYQS5YkLWSspXeTrGEQ6B+uqk/NPT8c8lW1Kcn7kqytqvsmV+qR55+5k7TSjXP3S4DrgJ1V9a4F2jypa0eSs7v3vX+ShUqSRhtnpP584BXA3Um2dsfeDJwCUFXXAhcBr0lyAPgBcHFV1eTLlSQtZmSoV9WXgIxocw1wzaSKkiQdmqb/nJ1/wk7SatN0qI/D4JfUEtd+kaSGGOqS1BBDXZIaYqhLUkMMdUlqiKEuSQ0x1CWpIav+PvWFuLiXpJXIkbokNcRQl6SGGOqS1BBDXZIaYqhLUkMMdUlqiKEuSQ0x1CWpIYa6JDXEUJekhhjqktQQQ12SGmKoS1JDXKVRE+GqltJ0cKQuSQ0x1CWpIYa6JDXEUJekhhjqktSQkaGe5OQktyTZkWR7ktfN0yZJ3p1kV5K7kpx1ZMqVJC1mnFsaDwBvqKo7khwHbElyc1XtGGrzYuD07vUc4P3dV0nSMho5Uq+qvVV1R7f9PWAncOKcZhcCN9bAl4Hjk5ww8WolSYta0px6klng2cBtc06dCNwztL+HxwY/SdYn2Zxk8/79+5dYqiRplLFDPcmxwCeB11fVw4fyYVW1oarWVdW6mZmZQ3kLSdIixgr1JGsYBPqHq+pT8zS5Fzh5aP+k7pgkaRmNc/dLgOuAnVX1rgWabQQu7e6CeS7wUFXtnWCdkqQxjHP3y/OBVwB3J9naHXszcApAVV0LbALOB3YB3wdePfFKJUkjjQz1qvoSkBFtCnjtpIqSJB0anyiVpIYY6pLUEENdkhpiqEtSQwx1SWqIoS5JDTHUJakh4zx8tOrNXvHZR7d3X31Bj5VI0uIcqUtSQwx1SWqIoS5JDTHUJakhhrokNcRQl6SGGOqS1BBDXZIaYqhLUkMMdUlqiKEuSQ0x1CWpIYa6JDXEVRoPg6s3Spo2zYX6cNBK0mrj9IskNcRQl6SGGOqS1JDm5tSnwaQuoHohVtJSGepLtNCFWC/QSpoGTr9IUkNGhnqS65PsS7JtgfPnJnkoydbuddXky5QkjWOc6ZcbgGuAGxdp88WqeslEKpIkHbKRI/WquhV4YBlqkSQdpknNqT8vyZ1JPpfkjIUaJVmfZHOSzfv375/QR0uSDppEqN8BnFpVzwLeA3x6oYZVtaGq1lXVupmZmQl8tCRp2GGHelU9XFWPdNubgDVJ1h52ZZKkJTvsUE/ypCTpts/u3vP+w31fSdLSjbz7JclNwLnA2iR7gLcAawCq6lrgIuA1SQ4APwAurqo6YhVLkhY0MtSr6pIR569hcMujJKlnPlEqSQ0x1CWpIYa6JDXEUJekhhjqktQQQ12SGmKoS1JDDHVJaoihLkkNMdQlqSGGuiQ1xFCXpIaM8zdKNQVmr/jso9u7r76gx0okTTNH6pLUEENdkhpiqEtSQwx1SWqIF0p74oVPSUeCI3VJaogj9WU0PDqXpCPBkbokNcRQl6SGGOqS1BBDXZIa4oXSI8yLo5KWkyN1SWqIoS5JDTHUJakhhrokNWRkqCe5Psm+JNsWOJ8k706yK8ldSc6afJmSpHGMM1K/AThvkfMvBk7vXuuB9x9+WZKkQzEy1KvqVuCBRZpcCNxYA18Gjk9ywqQKlCSNbxJz6icC9wzt7+mOPUaS9Uk2J9m8f//+CXy0JGnYsl4oraoNVbWuqtbNzMws50dL0qowiVC/Fzh5aP+k7pgkaZlNItQ3Apd2d8E8F3ioqvZO4H0lSUs0cu2XJDcB5wJrk+wB3gKsAaiqa4FNwPnALuD7wKuPVLGSpMWNDPWqumTE+QJeO7GKJEmHzCdKJakhhrokNcRQl6SGGOqS1BBDXZIaYqhLUkP8G6WauOG/y7r76gt6rERafRypS1JDHKlPAUe2kibFkbokNcRQl6SGGOqS1BBDXZIa4oVSaRXxonz7HKlLUkMMdUlqiKEuSQ0x1CWpIYa6JDXEUJekhhjqktQQQ11arc44Ax56qO8qNGE+fKRDNvwgi1agHTtg0ya45JK+K9EEOVKXVrNLL4Vjj4WXvazvSjQhjtRXCR8P17zWrIFTT4W3v73vSjQhhrq0Ss2+6TOPbu8+7bQeK9EkGeqrkKN2qV3OqUtSQwx1SWrIWKGe5LwkX0uyK8kV85x/VZL9SbZ2r9+bfKmSpFFGzqknOQp4L/AiYA9we5KNVbVjTtOPVdXlR6DGRXmvtCT9xDgj9bOBXVX1rar6IfBR4MIjW5Yk6VCMc/fLicA9Q/t7gOfM0+63kpwDfB34o6q6Z26DJOuB9QCnnHLK0qvV2PwNRlqdJnWh9J+B2ap6JnAz8MH5GlXVhqpaV1XrZmZmJvTRkqSDxhmp3wucPLR/UnfsUVV1/9DuB4B3Hn5pkibB39pWl3FG6rcDpyd5cpKjgYuBjcMNkpwwtPtSYOfkSpQkjWvkSL2qDiS5HPg8cBRwfVVtT/I2YHNVbQT+MMlLgQPAA8CrjmDNq55PhEpayFjLBFTVJmDTnGNXDW1fCVw52dIkSUvlE6WS1BAX9GqIF8QkGeqSHjMg8FrNyuX0iyQ1xJH6CueUi6RhjtQlqSGGuiQ1xFCXpIY4p77K+XSq1BZH6pLUEENdkhri9MuU8RZFSYfDkbokNcRQl6SGGOqS1BBDXZIaYqhLUkMMdUlqiKEuSQ0x1CWpIT58pEe5Doy08jlSl6SGGOqS1BBDXZIa4py6pMfw+srKZahrXv5Qr2yu9rl6Of0iSQ1xpK4lcQQoTTdH6pLUkLFCPcl5Sb6WZFeSK+Y5f0ySj3Xnb0syO/FKJUkjjZx+SXIU8F7gRcAe4PYkG6tqx1Czy4DvVtUvJrkYeAfwu0eiYHAKYLnZ3yuD/06C8ebUzwZ2VdW3AJJ8FLgQGA71C4G3dtv/AFyTJFVVE6xVK9BCQeMdNZOxHEHunVAryzihfiJwz9D+HuA5C7WpqgNJHgKeANw33CjJemB9t/tIkq9122vntp1C1jgZa4H78o6+y1jUtPdjb/Ut4d9t2vsQVm6Npy72Dct690tVbQA2zD2eZHNVrVvOWpbKGifDGg/ftNcH1jgph1LjOBdK7wVOHto/qTs2b5skjwN+Drh/KYVIkg7fOKF+O3B6kicnORq4GNg4p81G4JXd9kXAF5xPl6TlN3L6pZsjvxz4PHAUcH1VbU/yNmBzVW0ErgM+lGQX8ACD4F+Kx0zJTCFrnAxrPHzTXh9Y46QsucY4oJakdvhEqSQ1xFCXpIb0HuqjliCYBkl2J7k7ydYkm/uuByDJ9Un2Jdk2dOzxSW5O8o3u689PYY1vTXJv15dbk5zfY30nJ7klyY4k25O8rjs+Nf24SI3T1I8/neQrSe7savzz7viTu2VDdnXLiBw9hTXekOTbQ/14Zl81dvUcleSrST7T7S+9D6uqtxeDC6/fBJ4CHA3cCTy9z5oWqHM3sLbvOubUdA5wFrBt6Ng7gSu67SuAd0xhjW8F3th3/3W1nACc1W0fB3wdePo09eMiNU5TPwY4ttteA9wGPBf4OHBxd/xa4DVTWOMNwEV99+FQnX8MfAT4TLe/5D7se6T+6BIEVfVD4OASBBqhqm5lcKfRsAuBD3bbHwR+YzlrmmuBGqdGVe2tqju67e8BOxk8HT01/bhIjVOjBh7pdtd0rwJeyGDZEOi/HxeqcWokOQm4APhAtx8OoQ/7DvX5liCYqv9gOwX8a5It3VIH0+qJVbW32/4f4Il9FrOIy5Pc1U3P9DpFdFC3suizGYzgprIf59QIU9SP3bTBVmAfcDOD38AfrKoDXZPef7bn1lhVB/vxL7p+/Kskx/RXIX8N/Cnw427/CRxCH/Yd6ivFC6rqLODFwGuTnNN3QaPU4Pe1qRqJdN4PnAacCewF/rLXaoAkxwKfBF5fVQ8Pn5uWfpynxqnqx6r6v6o6k8ET52cDT+uznvnMrTHJM4ArGdT6q8DjgTf1UVuSlwD7qmrL4b5X36E+zhIEvauqe7uv+4B/ZPAf7TT6TpITALqv+3qu5zGq6jvdD9ePgb+l575MsoZBWH64qj7VHZ6qfpyvxmnrx4Oq6kHgFuB5wPHdsiEwRT/bQzWe101vVVX9L/B39NePzwdemmQ3g2noFwJ/wyH0Yd+hPs4SBL1K8rNJjju4Dfw6sG3x7+rN8HINrwT+qcda5nUwLDu/SY992c1ZXgfsrKp3DZ2amn5cqMYp68eZJMd32z/D4G8v7GQQnBd1zfrux/lq/K+h/3mHwXx1L/1YVVdW1UlVNcsgB79QVS/nUPpwCq72ns/giv43gT/ru5556nsKg7ty7gS2T0uNwE0Mfu3+EYO5tssYzMH9O/AN4N+Ax09hjR8C7gbuYhCeJ/RY3wsYTK3cBWztXudPUz8uUuM09eMzga92tWwDruqOPwX4CrAL+ARwzBTW+IWuH7cBf093h0yfL+BcfnL3y5L70GUCJKkhfU+/SJImyFCXpIYY6pLUEENdkhpiqEtSQwx1SWqIoS5JDfl/Yqu69Z0o4zsAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.hist(CH_convertion_to_pm, bins=100)\n",
    "ax.scatter(mean_converion_rate, 6e5, marker=\"*\", c='r')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6db22898",
   "metadata": {},
   "source": [
    "We can see a clear separation between the two groups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "67fd8119",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXoAAAEDCAYAAAA7jc+ZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAPuklEQVR4nO3df5BdZ13H8feHpC0OBao2nek00daaUjMIBdYCIlIRtD+cRuWHDTIW6RAHaIcfwkyqTK3tH7agjjgTxDBWlBFKQAbjNBpFijBMW7IRKCSlGEqxWyoNpQUZBkrl6x/3xFyWzd67m/tj8+T9mtnJPc959p7vPrn7uc99zrl3U1VIktr1qGkXIEkaL4Nekhpn0EtS4wx6SWqcQS9JjTPoJalxUw36JDckuT/JZ4fs/+Ik+5LsTfLucdcnSS3INK+jT/LzwDeBv62qJw7oux7YDjy3qh5MckpV3T+JOiXpaDbVGX1VfRT4Wn9bkjOT/HOSPUk+luTsbtcrgK1V9WD3vYa8JA1hJa7RbwOuqKqnAW8A3ta1nwWcleTjSW5Ncv7UKpSko8jqaRfQL8mJwM8C70tysPmE7t/VwHrgPGAt8NEkP11VD024TEk6qqyooKf3CuOhqjpngX1zwG1V9V3gi0k+Ty/4d0+wPkk66qyopZuq+ga9EH8RQHqe3O3+IL3ZPElOpreUc9cUypSko8rAoB90CWQXxn+eZH+S25M8ddiDJ3kPcAvwhCRzSS4DfhO4LMmngb3Axq77LuCBJPuAm4E3VtUDwx5Lko5VAy+vHHQJZJILgSuAC4GnA2+tqqePoVZJ0jIMnNEvdAnkPBvpPQlUVd0KnJTk1FEVKEk6MqM4GXsacE/f9lzXdt/8jkk2A5sBHvOYxzzt7LPPnt9FkrSIPXv2fLWq1izleyZ61U1VbaN3nTwzMzM1Ozs7ycNL0lEvyZeW+j2juOrmXmBd3/bark2StAKMIuh3AL/VXX3zDODrVfUDyzaSpOkYuHTTXQJ5HnBykjngD4DjAKrq7cBOelfc7Ae+Bfz2uIqVJC3dwKCvqk0D9hfw6pFVJEkaqRX1zlhJ0ugZ9JLUOINekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuMMeklqnEEvSY0z6CWpcQa9JDXOoJekxhn0ktQ4g16SGmfQS1LjDHpJapxBL0mNM+glqXEGvSQ1zqCXpMYZ9JLUOINekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaN1TQJzk/yZ1J9ifZssD+H0tyc5JPJrk9yYWjL1WStBwDgz7JKmArcAGwAdiUZMO8bm8CtlfVU4BLgLeNulBJ0vIMM6M/F9hfVXdV1cPAjcDGeX0KeFx3+/HAl0dXoiTpSAwT9KcB9/Rtz3Vt/a4GXppkDtgJXLHQHSXZnGQ2yeyBAweWUa4kaalGdTJ2E/DOqloLXAi8K8kP3HdVbauqmaqaWbNmzYgOLUlazDBBfy+wrm97bdfW7zJgO0BV3QI8Gjh5FAVKko7MMEG/G1if5Iwkx9M72bpjXp//An4RIMlP0Qt612YkaQUYGPRV9QhwObALuIPe1TV7k1yT5OKu2+8Cr0jyaeA9wMuqqsZVtCRpeKuH6VRVO+mdZO1vu6rv9j7gWaMtTZI0Cr4zVpIaZ9BLUuMMeklqnEEvSY0b6mSstFynb7lp7Me4+7qLxn4M6WjmjF6SGmfQS1LjDHpJapxBL0mNM+glqXEGvSQ1zqCXpMYZ9JLUOINekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuMMeklqnEEvSY0z6CWpcQa9JDXOoJekxhn0ktQ4g16SGmfQS1LjDHpJapxBL0mNM+glqXFDBX2S85PcmWR/ki2H6fPiJPuS7E3y7tGWKUlartWDOiRZBWwFng/MAbuT7KiqfX191gNXAs+qqgeTnDKugiVJSzPMjP5cYH9V3VVVDwM3Ahvn9XkFsLWqHgSoqvtHW6a0iJe8ZNoVSCvaMEF/GnBP3/Zc19bvLOCsJB9PcmuS8xe6oySbk8wmmT1w4MDyKpbmu/baaVcgrWijOhm7GlgPnAdsAt6R5KT5napqW1XNVNXMmjVrRnRoHfPOPHPaFUgr2jBBfy+wrm97bdfWbw7YUVXfraovAp+nF/ySpCkbeDIW2A2sT3IGvYC/BJi/KPpBejP5v05yMr2lnLtGWKd0WKdvuWnsx7j7uovGfgxpXAbO6KvqEeByYBdwB7C9qvYmuSbJxV23XcADSfYBNwNvrKoHxlW0JGl4w8zoqaqdwM55bVf13S7g9d2XJGkF8Z2xktQ4g16SGmfQS1LjDHpJapxBL0mNM+glqXEGvSQ1zqCXpMYZ9JLUOINekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuMMeklqnEEvSY0z6CWpcaunXYCm5/QtN027BEkT4Ixekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuMMeklqnEEvSY0bKuiTnJ/kziT7k2xZpN8LklSSmdGVKEk6EgODPskqYCtwAbAB2JRkwwL9Hgu8Brht1EVKkpZvmBn9ucD+qrqrqh4GbgQ2LtDvWuB64NsjrE+SdISGCfrTgHv6tue6tv+X5KnAuqpa9APOk2xOMptk9sCBA0suVpK0dEf8h0eSPAr4U+Blg/pW1TZgG8DMzEwd6bGlSZnEH2m5+7qLxn4MHZuGmdHfC6zr217btR30WOCJwEeS3A08A9jhCVlJWhmGCfrdwPokZyQ5HrgE2HFwZ1V9vapOrqrTq+p04Fbg4qqaHUvFkqQlGRj0VfUIcDmwC7gD2F5Ve5Nck+TicRcoSToyQ63RV9VOYOe8tqsO0/e8Iy9LkjQqvjNWkhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuMMeklqnEEvSY0z6CWpcQa9JDXOoJekxhn0ktQ4g16SGmfQS1LjDHpJapxBL0mNM+glqXEGvSQ1zqCXpMYZ9JLUOINekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuOGCvok5ye5M8n+JFsW2P/6JPuS3J7k35L8+OhLlSQtx8CgT7IK2ApcAGwANiXZMK/bJ4GZqnoS8H7gzaMuVJK0PMPM6M8F9lfVXVX1MHAjsLG/Q1XdXFXf6jZvBdaOtkxJ0nINE/SnAff0bc91bYdzGfBPC+1IsjnJbJLZAwcODF+lJGnZRnoyNslLgRngLQvtr6ptVTVTVTNr1qwZ5aElSYexeog+9wLr+rbXdm3fJ8nzgN8HnlNV3xlNeZKkIzXMjH43sD7JGUmOBy4BdvR3SPIU4C+Bi6vq/tGXKUlaroFBX1WPAJcDu4A7gO1VtTfJNUku7rq9BTgReF+STyXZcZi7kyRN2DBLN1TVTmDnvLar+m4/b8R1SZJGxHfGSlLjDHpJapxBL0mNM+glqXEGvSQ1zqCXpMYZ9JLUOINekhpn0EtS4wx6SWqcQS9JjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuOG+gtTksbv9C03jf0Yd1930diPoZXHGb0kNc6gl6TGuXSzQk3iZbykY4MzeklqnEEvSY0z6CWpcQa9JDXOoJekxhn0ktQ4g16SGmfQS1LjDHpJapxBL0mNM+glqXEGvSQ1zg81WwY/cEzS0cQZvSQ1bqigT3J+kjuT7E+yZYH9JyR5b7f/tiSnj7xSSdKyDFy6SbIK2Ao8H5gDdifZUVX7+rpdBjxYVT+Z5BLgeuA3xlHwIC6rSIfnnys8NqWqFu+QPBO4uqp+udu+EqCq/qivz66uzy1JVgP/DaypRe78hFPX16mX/tmR/wSSjjnH8pNJkj1VNbOU7xnmZOxpwD1923PA0w/Xp6oeSfJ14EeBr84rcDOwudv8zpeu/5XPLqXYhp3MvLE6hjkWhzgWh3zfWOT6KVYyfU9Y6jdM9KqbqtoGbANIMrvUZ6VWORaHOBaHOBaHOBaHJJld6vcMczL2XmBd3/barm3BPt3SzeOBB5ZajCRp9IYJ+t3A+iRnJDkeuATYMa/PDuDS7vYLgQ8vtj4vSZqcgUs33Zr75cAuYBVwQ1XtTXINMFtVO4C/At6VZD/wNXpPBoNsO4K6W+NYHOJYHOJYHOJYHLLksRh41Y0k6ejmO2MlqXEGvSQ1buxBn2RdkpuT7EuyN8lruvark9yb5FPd14XjrmXaDjcW3b4rknyua3/zNOuchEUeF+/te0zcneRTUy517BYZi3OS3NqNxWySc6dd67gtMhZPTnJLks8k+cckj5t2reOW5NFJPpHk091Y/GHXfkb3UTP7u9+X4wfeWVWN9Qs4FXhqd/uxwOeBDcDVwBvGffyV9LXIWPwC8CHghG7fKdOudVpjMa/PnwBXTbvWKT4u/gW4oGu/EPjItGud4ljsBp7Ttb8cuHbatU5gLAKc2N0+DrgNeAawHbika3878MpB9zX2GX1V3VdV/9Hd/h/gDnrvpD3mLDIWrwSuq6rvdPvun16VkzHocZEkwIuB90ynwslZZCwKODhzfTzw5elUODmLjMVZwEe7bv8KvGA6FU5O9Xyz2zyu+yrgucD7u/a/AX510H1NdI2++1TLp9B7ZgK4PMntSW5I8sOTrGXa5o3FWcCzu5dj/57kZ6Za3IQt8LgAeDbwlar6z6kUNSXzxuK1wFuS3AP8MXDl9CqbvHljsRfY2O16Ed//Js5mJVnVLV/eT+8J7gvAQ1X1SNdljiEmzhML+iQnAn8PvLaqvgH8BXAmcA5wH72X6ceEBcZiNfAj9F6WvRHY3s1om7fAWBy0iWNgNt9vgbF4JfC6qloHvI7e+1WOCQuMxcuBVyXZQ29J5+Fp1jcpVfW/VXUOvU8kOBc4ezn3M5GgT3Icvf+0v6uqDwBU1Ve6H+J7wDvo/RDNW2gs6D0rf6B7qfYJ4Hv0PsSpaYcZi4Mfo/HrwHunVdukHWYsLgUO3n4fx/DvSFV9rqp+qaqeRm8C8IVp1jhpVfUQcDPwTOCk7ncEFv5Imh8wiatuQm8mckdV/Wlf+6l93X4NaP6TLA83FsAH6Z2QJclZwPE0/qmFi4wFwPOAz1XV3OQrm7xFxuLLwHO6288Fml/GWiQvTun+fRTwJnonIZuWZE2Sk7rbP0Tvb4LcQS/wX9h1uxT4h4H31Z25HZskPwd8DPgMvZkqwO/Re2l+Dr2TC3cDv1NV9421mClbZCw+BNxAbzwepnc10oenUeOkHG4sqmpnkncCt1ZV87/MsOjj4hvAW+kt7X0beFVV7ZlKkROyyFisB17dbX8AuLLGHV5TluRJ9E62rqI3Kd9eVdck+QngRnrLvZ8EXnrwQo7D3lfjYyVJxzzfGStJjTPoJalxBr0kNc6gl6TGGfSS1DiDXpIaZ9BLUuP+D9QwCYv6pX5rAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.hist(CH_convertion_to_pm, bins=100)\n",
    "ax.scatter(mean_converion_rate, 6e5, marker=\"*\", c='r')\n",
    "\n",
    "ax.set_xlim(25, 30)\n",
    "ax.set_ylim(0, 1e6)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d54bef37",
   "metadata": {},
   "source": [
    "The convertion rate is around $27.15$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e91d299b",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONVERTION_RATE = 27.15"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6e4ba1d",
   "metadata": {},
   "source": [
    "## Is there a bond between atoms?\n",
    "\n",
    "For each bond there are approximate mean lengths, however, that is not enough to assess whether there is a bond between atoms. A better way to tell is to look at if the bond length is in a certain range. \n",
    "\n",
    "Let's figure out all atom pairs there are and look at how distances between them are distributed. It is expected that there would be a division somewhere in the range $[100, 250]$ because that is where the pair-covalent radius sums lie."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "433ab519",
   "metadata": {},
   "outputs": [],
   "source": [
    "# record all atom pair distances\n",
    "\n",
    "bond_lengths = {}\n",
    "\n",
    "for i in tqdm(range(1, TOTAL_MOLECULES+1)):\n",
    "    row = train_db.get(i)\n",
    "    symbols = row.symbols\n",
    "    positions = row.positions\n",
    "\n",
    "    distance_matrix = position_to_distance_matrix(positions)\n",
    "\n",
    "    for i, a1 in enumerate(symbols):\n",
    "        for j, a2 in enumerate(symbols):\n",
    "            if i <= j: continue\n",
    "                \n",
    "            bond_name = \"\".join(sorted([a1, a2]))\n",
    "            bond_length = distance_matrix[i][j] * CONVERTION_RATE\n",
    "\n",
    "            if bond_length > 250 or bond_length < 100: continue\n",
    "            \n",
    "            if bond_name not in bond_lengths:\n",
    "                bond_lengths[bond_name] = [bond_length]\n",
    "            else:\n",
    "                bond_lengths[bond_name].append(bond_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98c60d75",
   "metadata": {},
   "outputs": [],
   "source": [
    "for bond_name in bond_lengths:\n",
    "    bond_lengths[bond_name] = np.array(bond_lengths[bond_name])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63b8fa1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(bond_lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f15a3d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(6, 6, figsize=(20, 10))\n",
    "\n",
    "for ii, bond_name in enumerate(tqdm(sorted(bond_lengths, key=lambda x: len(bond_lengths[x]), reverse=True))):\n",
    "    i=ii//6\n",
    "    j=ii% 6\n",
    "    ax = axs[i][j]\n",
    "    ax.set_xlim([100, 250])\n",
    "    ax.title.set_text(bond_name)\n",
    "    \n",
    "    amount, buckets, _ = ax.hist(bond_lengths[bond_name], bins=50)\n",
    "\n",
    "    zero_x = buckets[np.where(amount <= 5)[0]][-1]\n",
    "    ax.text(zero_x, np.zeros_like(zero_x), str(round(zero_x)), rotation=45, c='r')\n",
    "    \n",
    "plt.subplots_adjust(\n",
    "    left=0.1,\n",
    "    bottom=0.1, \n",
    "    right=1, \n",
    "    top=1, \n",
    "    wspace=0.5, \n",
    "    hspace=0.5\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec695c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_bond_length = {\n",
    "    'CH':163,\n",
    "    'CC':189,\n",
    "    'CN':188,\n",
    "    'HN':150,\n",
    "    'CO':189,\n",
    "    'HO':120,\n",
    "    'NO':189,\n",
    "    'NN':182,\n",
    "    'CS':200,\n",
    "    'CF':210,\n",
    "    'OS':200,\n",
    "    'NS':203,\n",
    "    'CCl':240,\n",
    "    'BrC':230,\n",
    "    'BrN':200\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f317b4a5",
   "metadata": {},
   "source": [
    "## Getting the bond type matrix\n",
    "\n",
    "Now we can try to get the bond type matrix. There are two ways we can achive that:\n",
    "1. Compare the distance between atoms to the average length of each bond type\n",
    "2. Make a CSP that assigns the bond types that satisfy the _bond constraint*_\n",
    "\n",
    "The problem with the first method is that it doesn't account for the bond constraint. The problem with the second method is that there could be multiple ways atoms are bonded and hence we have to choose a csp solution somehow.\n",
    "\n",
    "To cancel the both problems we will find solutions using a CSP search and if there are multiple solutions then we will look at the one that is the \"closest\" to the solution from the first method.\n",
    "\n",
    "\\* – bond happends between two atoms when they share electrons. Each atom has a total amount of electrons it can give/take. A bond type is the number of electrons atoms share in a bond."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34d123e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get a list of all atoms in the dataset\n",
    "\n",
    "atoms_in_dataset = set()\n",
    "\n",
    "for i in tqdm(range(1, TOTAL_MOLECULES + 1)):\n",
    "    row = train_db.get(i)\n",
    "    atoms = row.symbols\n",
    "\n",
    "    atoms_in_dataset.update(atoms)\n",
    "\n",
    "atoms_in_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e57dd5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# molecule to test the functions\n",
    "\n",
    "row = train_db.get(64)\n",
    "atom_names = row.symbols\n",
    "positions = row.positions\n",
    "\n",
    "distance_matrix = position_to_distance_matrix(positions) * CONVERTION_RATE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92ac3fb2",
   "metadata": {},
   "source": [
    "### Distance to (possible) bond matrix\n",
    "\n",
    "For this we will use the maximum bond lengths found earlier. Some atoms are not connected though. That's why it's \"possible\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "331cd42e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def possibly_is_bond(atom_1, atom_2, bond_length):\n",
    "    bond_name = \"\".join(sorted([atom_1, atom_2]))\n",
    "\n",
    "    if (bond_name in max_bond_length) and (bond_length < max_bond_length[bond_name]):\n",
    "        return 1\n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "# test the function\n",
    "possibly_is_bond(atom_names[2], atom_names[8], distance_matrix[2][8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61d585ec",
   "metadata": {},
   "source": [
    "Looks good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e94ebd71",
   "metadata": {},
   "outputs": [],
   "source": [
    "def distance_to_possible_bond_matrix(distance_matrix, atoms):\n",
    "    n_atoms = len(atoms)\n",
    "    \n",
    "    bond_matrix = np.zeros_like(distance_matrix)\n",
    "    \n",
    "    for i in range(n_atoms-1):\n",
    "        for j in range(i+1, n_atoms):\n",
    "            is_bond = possibly_is_bond(atoms[i], atoms[j], distance_matrix[i][j])\n",
    "            \n",
    "            bond_matrix[i][j] = is_bond\n",
    "            bond_matrix[j][i] = is_bond\n",
    "\n",
    "    return bond_matrix\n",
    "\n",
    "# test the function\n",
    "possible_bond_matrix = distance_to_possible_bond_matrix(distance_matrix, atom_names)\n",
    "\n",
    "# number of possible bonds each atom has\n",
    "possible_bond_matrix.sum(axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3146fe93",
   "metadata": {},
   "source": [
    "From here errors are already seen. Some atoms have too many bonds..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad6a9abf",
   "metadata": {},
   "source": [
    "### Visualize the molecule [bonus]\n",
    "\n",
    "This is not needed for the problem, however it is useful to be able to inspect the molecule if any errors arise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7b8a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def str_int(n):\n",
    "    if type(n) is float or type(n) is int:\n",
    "        return int(n) if int(n) == n else n\n",
    "    else:\n",
    "        return n\n",
    "\n",
    "def display_molecule_video(row, bond_matrix, atom_text=None, bond_text=None, equal_aspect_ratio=True, angle=10):\n",
    "    fig = plt.figure()\n",
    "    fig.set_size_inches(12, 8)\n",
    "\n",
    "    ax = plt.axes(projection='3d')\n",
    "    if equal_aspect_ratio:\n",
    "        p = row.positions\n",
    "        p_max = p.max()\n",
    "        p_min = p.min()\n",
    "\n",
    "        ax.set_xlim([p_min, p_max])\n",
    "        ax.set_ylim([p_min, p_max])\n",
    "        ax.set_zlim([p_min, p_max])\n",
    "    # ax.grid(False)\n",
    "    \n",
    "    def init():\n",
    "        p = row.positions\n",
    "        s = row.symbols\n",
    "        n = row.numbers\n",
    "        \n",
    "        x = p[:, 0]\n",
    "        y = p[:, 1]\n",
    "        z = p[:, 2]\n",
    "\n",
    "        ax.scatter(x, y, z, \n",
    "            s = n*20,        # size\n",
    "            c = n,           # color\n",
    "            cmap = 'inferno'\n",
    "        )\n",
    "        \n",
    "        if type(atom_text) != type(None):\n",
    "            for i in range(len(p)):\n",
    "                ax.text(x[i], y[i], z[i], str(atom_text[i]))\n",
    "\n",
    "        for i in range(len(p)):\n",
    "            for j in range(i+1, len(p)):\n",
    "                if bond_matrix[i][j] != 0:\n",
    "                    lx = np.array([x[i], x[j]])\n",
    "                    ly = np.array([y[i], y[j]])        \n",
    "                    lz = np.array([z[i], z[j]])\n",
    "\n",
    "                    ax.plot(lx, ly, lz, color='black')\n",
    "                    if type(bond_text) != type(None):\n",
    "                        ax.text(lx.mean(), ly.mean(), lz.mean(), str_int(bond_text[i][j]))\n",
    "        return fig,\n",
    "\n",
    "\n",
    "    def animate(i):\n",
    "        ax.view_init(elev=angle, azim=i)\n",
    "        return fig,\n",
    "\n",
    "    \n",
    "    anim = animation.FuncAnimation(fig, animate, init_func=init,\n",
    "                                   frames=360, interval=30, blit=True)\n",
    "    \n",
    "    plt.close()\n",
    "    return HTML(anim.to_html5_video())\n",
    "\n",
    "# test the function\n",
    "display_molecule_video(\n",
    "    row, \n",
    "    possible_bond_matrix, \n",
    "#     atom_text=list(range(row.natoms)),\n",
    "    atom_text=atom_names, \n",
    "    bond_text=None, \n",
    "    angle=40\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "737efbce",
   "metadata": {},
   "source": [
    "## CSP for molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71ada931",
   "metadata": {},
   "outputs": [],
   "source": [
    "from funcs.csp import Problem as CSP\n",
    "from funcs.csp import ExactSumConstraint, BacktrackingSolver\n",
    "\n",
    "total_bonds = {\n",
    "    'Br': 1, \n",
    "    'C' : 4, \n",
    "    'Cl': 1, \n",
    "    'F' : 1, \n",
    "    'H' : 1, \n",
    "    'N' : 3, \n",
    "    'O' : 2, \n",
    "    'S' : [2, 4, 6]\n",
    "}\n",
    "\n",
    "def get_bond_constraint(total_bonds):\n",
    "    def bond_constraint(*bonds):\n",
    "        bond_sum = sum(bonds)\n",
    "\n",
    "        if type(total_bonds) is list:\n",
    "            return bond_sum in total_bonds\n",
    "        else:\n",
    "            return bond_sum == total_bonds\n",
    "\n",
    "    return bond_constraint\n",
    "\n",
    "def sort_tuple(t):\n",
    "    return tuple(sorted(t))\n",
    "\n",
    "def csp_solve_bond_matrix(row, bond_matrix, get_bond_names=False, all_sols=True):\n",
    "    symbols = row.symbols\n",
    "    n_atoms = row.natoms\n",
    "    \n",
    "    variables = []\n",
    "    bond_names = {}\n",
    "    \n",
    "    for i in range(n_atoms-1):\n",
    "        for j in range(i+1, n_atoms):\n",
    "            if bond_matrix[i][j]:\n",
    "                variables.append(sort_tuple((i, j)))\n",
    "    \n",
    "    csp = CSP(BacktrackingSolver())\n",
    "    \n",
    "    for variable in variables:\n",
    "        i, j = variable\n",
    "        a1 = symbols[i]\n",
    "        a2 = symbols[j]\n",
    "        bond_name = \"\".join(sorted([a1, a2]))\n",
    "        if bond_name not in bond_names:\n",
    "            bond_names[bond_name] = [variable]\n",
    "        else:\n",
    "            bond_names[bond_name].append(variable)\n",
    "        \n",
    "        if a1 == 'S':\n",
    "            number_of_bond_types = min(total_bonds[a1][0], total_bonds[a2])\n",
    "        elif a2 == 'S':\n",
    "            number_of_bond_types = min(total_bonds[a1], total_bonds[a2][0])\n",
    "        else:\n",
    "            number_of_bond_types = min(total_bonds[a1], total_bonds[a2])\n",
    "            \n",
    "        domain = list(range(1, number_of_bond_types+1))\n",
    "        \n",
    "        if bond_name == 'CC' or bond_name == 'CN':\n",
    "            domain = [1, 2, 3, 1.5]\n",
    "        \n",
    "        csp.addVariable(variable, domain)\n",
    "\n",
    "    for i in range(n_atoms):\n",
    "        a1 = symbols[i]\n",
    "        bonds = [sort_tuple((i, j)) for j in np.where(bond_matrix[i] > 0)[0]]\n",
    "        csp.addConstraint(\n",
    "            get_bond_constraint(total_bonds[a1]),\n",
    "            bonds\n",
    "        )\n",
    "            \n",
    "\n",
    "    if all_sols:\n",
    "        solutions = csp.getSolutions()\n",
    "    else:\n",
    "        solutions = csp.getSolution()\n",
    "        if solutions == None:\n",
    "            solutions = []\n",
    "        else:\n",
    "            solutions = [solutions]\n",
    "\n",
    "    if get_bond_names:\n",
    "        return bond_names, solutions\n",
    "    else:\n",
    "        return solutions\n",
    "\n",
    "csp_solutions = csp_solve_bond_matrix(row, possible_bond_matrix)\n",
    "\n",
    "print('Number of solutions:', len(csp_solutions))\n",
    "# print(csp_solutions)\n",
    "\n",
    "for sol in csp_solutions:\n",
    "    print('# of aromatic bonds:', list(sol.values()).count(1.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63bfea0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_csp_bond_type_matrix(csp_solution, bond_matrix):\n",
    "    bond_type_matrix = np.zeros_like(bond_matrix)\n",
    "\n",
    "    for i, j in csp_solution:\n",
    "        bond_type_matrix[i][j] = csp_solution[(i, j)]\n",
    "        bond_type_matrix[j][i] = csp_solution[(i, j)]\n",
    "        \n",
    "    return bond_type_matrix\n",
    "\n",
    "csp_bond_type_matrix = get_csp_bond_type_matrix(csp_solutions[0], possible_bond_matrix)\n",
    "print(csp_bond_type_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faaf36a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "csp_bond_matrix = (csp_bond_type_matrix != 0).astype('int')\n",
    "\n",
    "display_molecule_video(\n",
    "    row, \n",
    "    csp_bond_matrix,\n",
    "    atom_text=atom_names, \n",
    "    bond_text=csp_bond_type_matrix, \n",
    "    angle=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2cdfe77",
   "metadata": {},
   "outputs": [],
   "source": [
    "def csp_solution_to_bond_type_counts(bond_names, solution):\n",
    "    bond_type_counts = {}\n",
    "    \n",
    "    for bond_name in bond_names:\n",
    "        for var in bond_names[bond_name]:\n",
    "            if solution[var] > 0:\n",
    "                bond_type = bond_name + str(solution[var])\n",
    "                if bond_type in bond_type_counts:\n",
    "                    bond_type_counts[bond_type] += 1\n",
    "                else:\n",
    "                    bond_type_counts[bond_type] = 1\n",
    "                \n",
    "    return bond_type_counts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ada33f50",
   "metadata": {},
   "source": [
    "## Processing the whole dataset\n",
    "\n",
    "Now let's use all the functions to make the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "400bb7cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame()\n",
    "csp_failed_to_solve = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d8fb0c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in tqdm(failed):\n",
    "    row = train_db.get(i)\n",
    "    atom_names = row.symbols\n",
    "    positions = row.positions\n",
    "\n",
    "    distance_matrix = position_to_distance_matrix(positions) * CONVERTION_RATE\n",
    "    possible_bond_matrix = distance_to_possible_bond_matrix(distance_matrix, atom_names)\n",
    "    bond_names, csp_solutions = csp_solve_bond_matrix(row, possible_bond_matrix, True, all_sols=False)\n",
    "    \n",
    "    if len(csp_solutions) == 0:\n",
    "        csp_failed_to_solve.append(i)\n",
    "        continue\n",
    "    elif len(csp_solutions) == 1:\n",
    "        csp_solution = csp_solutions[0]\n",
    "        bond_type_counts = csp_solution_to_bond_type_counts(bond_names, csp_solution)\n",
    "    else:\n",
    "        # choose the one with the most aromatic bonds\n",
    "        csp_solutions.sort(key=lambda sol: list(sol.values()).count(1.5), reverse=True)\n",
    "        csp_solution = csp_solutions[0]\n",
    "        bond_type_counts = csp_solution_to_bond_type_counts(bond_names, csp_solution)\n",
    "\n",
    "    for bond in bond_type_counts:\n",
    "        if bond not in df.columns:\n",
    "            df[bond] = 0\n",
    "        df.loc[i, bond] = bond_type_counts[bond]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "644b76f5",
   "metadata": {},
   "source": [
    "### Making a sum of inverse distances dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ae7b76",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}