search_bonds.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a685664f",
   "metadata": {},
   "source": [
    "## Search bonds in molecular graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7a447106-7b38-4531-aec6-55e40b1e18e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import networkx as nx\n",
    "import ase\n",
    "import ase.visualize\n",
    "from scipy.spatial import cKDTree as KDTree\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a89606dd-ed0e-4c1b-85f5-8d5114942427",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = \"../champs-scalar-coupling/\"\n",
    "train = pd.read_csv(dataset + \"train.csv\")\n",
    "test = pd.read_csv(dataset + \"test.csv\")\n",
    "structures = pd.read_csv(dataset + \"structures.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "15316dc4-6ca7-48f8-a0d9-0ac93307ffda",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<html>\n",
       "\n",
       " <head>\n",
       "\n",
       "  <title>ASE atomic visualization</title>\n",
       "\n",
       "  <link rel=\"stylesheet\" type=\"text/css\"\n",
       "\n",
       "   href=\"https://www.x3dom.org/x3dom/release/x3dom.css\">\n",
       "\n",
       "  </link>\n",
       "\n",
       "  <script type=\"text/javascript\"\n",
       "\n",
       "   src=\"https://www.x3dom.org/x3dom/release/x3dom.js\">\n",
       "\n",
       "  </script>\n",
       "\n",
       " </head>\n",
       "\n",
       " <body>\n",
       "\n",
       "  <X3D>\n",
       "\n",
       "   <Scene>\n",
       "\n",
       "    <Transform translation=\"0.16 1.47 -0.07\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"0.565 0.565 0.565\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.76\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"-0.19 -0.01 0.12\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"0.565 0.565 0.565\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.76\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"0.49 -0.61 1.32\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"0.565 0.565 0.565\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.76\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"0.99 -1.98 1.35\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"0.565 0.565 0.565\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.76\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"1.89 -0.88 1.20\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 0.051 0.051\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.66\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"1.24 1.60 -0.18\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"-0.33 1.88 -0.96\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"-0.16 2.07 0.79\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"-1.28 -0.13 0.22\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"0.11 -0.58 -0.77\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"0.22 -0.15 2.28\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"1.06 -2.52 2.30\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "    <Transform translation=\"0.87 -2.62 0.47\">\n",
       "\n",
       "     <Shape>\n",
       "\n",
       "      <Appearance>\n",
       "\n",
       "       <Material diffuseColor=\"1.000 1.000 1.000\" specularColor=\"0.5 0.5 0.5\">\n",
       "\n",
       "       </Material>\n",
       "\n",
       "      </Appearance>\n",
       "\n",
       "      <Sphere radius=\"0.31\">\n",
       "\n",
       "      </Sphere>\n",
       "\n",
       "     </Shape>\n",
       "\n",
       "    </Transform>\n",
       "\n",
       "   </Scene>\n",
       "\n",
       "  </X3D>\n",
       "\n",
       " </body>\n",
       "\n",
       "</html>\n",
       "\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = structures[structures[\"molecule_name\"] == \"dsgdb9nsd_000150\"]\n",
    "ase.visualize.view(ase.Atoms(positions=a[['x','y','z']].values, symbols=a.atom.values), viewer=\"x3d\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "25801b31-3050-4fef-9a61-3ad79f8a0ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"usual\" valence of atoms\n",
    "VALENCE_MAX = {'C': 4, \n",
    "               'H': 1, \n",
    "               'N': 4, \n",
    "               'O': 2,\n",
    "               'F': 1}\n",
    "VALENCE_STD = {'C': 4,\n",
    "               'H': 1,\n",
    "               'N': 3,\n",
    "               'O': 2,\n",
    "               'F': 1}\n",
    "\n",
    "# expected distances in [A] for covalence 1 bond\n",
    "BOND_DIST_C1 = {'C': 0.77,\n",
    "                'H': 0.38,\n",
    "                'N': 0.75,\n",
    "                'O': 0.73,\n",
    "                'F': 0.71}\n",
    "\n",
    "# order used for finding bonds by atom type\n",
    "BOND_ORDER = {'H': 0,\n",
    "              'F': 0,\n",
    "              'O': 1,\n",
    "              'N': 2,\n",
    "              'C': 3}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "50352ef4-c4b9-4021-aac1-c27da77fb770",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_bond(n_avail, nbond, a0, a1, d1=None):\n",
    "    key = tuple(sorted((a0, a1)))\n",
    "    if key in nbond:\n",
    "        nbond[key][0] += 1.0\n",
    "    elif d1 is not None:\n",
    "        nbond[key] = [1.0, d1]\n",
    "    else:\n",
    "        raise Exception(f\"{a0},{a1} added after phase 1\")\n",
    "    n_avail[a0] -= 1\n",
    "    n_avail[a1] -= 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "927be37c-afeb-4df0-a42d-de9dc7597865",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bonded_atoms(atoms, nbond, i):\n",
    "    \"\"\"returns: [sorted atoms list], [sorted atom index] )\"\"\"\n",
    "    bonded = []\n",
    "    for (a0, a1), (n, _) in nbond.items():\n",
    "        if a0 == i:\n",
    "            bonded.append((a1, atoms[a1]))\n",
    "        elif a1 == i:\n",
    "            bonded.append((a0, atoms[a0]))\n",
    "    bonded = sorted(bonded, key=lambda b: b[1])\n",
    "    return \"\".join([b[1] for b in bonded]), [b[0] for b in bonded]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "88474ee6-67ce-4ebf-a6f4-15ac8c6bf549",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_bonds(kdt, n_avail, nbond, connected, isleaf, coords, atoms, atoms_idx, a0, connect_once=True, VALENCE=VALENCE_STD):\n",
    "    atom0 = atoms[a0]\n",
    "    if n_avail[a0] == 0:\n",
    "        return\n",
    "\n",
    "    # select closest atoms ORDERED BY DISTANCE: closest first\n",
    "    # note: the first answer is the atom itself and must be removed\n",
    "    next_dist, next_i = kdt.query(coords[a0], min(1+VALENCE[atom0], len(atoms)))\n",
    "    next_dist = next_dist[1:]; # remove a0 from list\n",
    "    next_i = next_i[1:]\n",
    "\n",
    "    # for each #VALENCE closest atoms\n",
    "    found = False\n",
    "    for d1, a1 in zip(next_dist, next_i):\n",
    "        if connect_once and (a1 in connected[a0]):\n",
    "            continue  # enforce 1-bond only in STEP 1\n",
    "        atom1 = atoms[a1]\n",
    "        predicted_bond = BOND_DIST_C1[atom0] + BOND_DIST_C1[atom1]\n",
    "        if abs(d1 / predicted_bond) < 1.2: # keep only atoms in the 20% expected distance or closer\n",
    "            if n_avail[a1] > 0:\n",
    "                add_bond(n_avail, nbond, a0, a1, d1)\n",
    "                connected[a0][a1] = 1\n",
    "                connected[a1][a0] = 1\n",
    "                if (n_avail[a0] == 0) or (n_avail[a1] == 0):\n",
    "                    isleaf[a0] = 1\n",
    "                    isleaf[a1] = 1\n",
    "                found = True\n",
    "                #print(\"leaF/Trunk & avail: \"+ \", \".join([f\"{i}:{atoms[i]}={leaflabel[isleaf[i]]}{n_avail[i]}\"\n",
    "                #                 for i in ordered_atoms_index]))\n",
    "\n",
    "        else:\n",
    "            #print(f\"-- match failure in molecule_name={name} {a0}:{atom0}-{a1}:{atoms[a1]}={d1} predicted={predicted_bond}\")\n",
    "            pass\n",
    "    return found"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3be72469-033d-43e4-9fc3-adc07d13aa64",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_bonds(structures, molecules):\n",
    "    out_name = []\n",
    "    out_a0 = []\n",
    "    out_a1 = []\n",
    "    out_n = []\n",
    "    out_dist = []\n",
    "    out_error = []\n",
    "    out_type = []\n",
    "    cycle_name = []\n",
    "    cycle_index = []\n",
    "    cycle_seq = []\n",
    "    cycle_atom_index = []\n",
    "    charge_name = []\n",
    "    charge_atom_index = []\n",
    "    charge_value = []\n",
    "    \n",
    "    for imol, name in tqdm(list(enumerate(molecules))):\n",
    "        molecule = structures.loc[name]\n",
    "        error    = 0\n",
    "        atoms    = molecule.atom.values\n",
    "        atoms_idx= molecule.atom_index.values\n",
    "        \n",
    "        n_avail  = np.asarray([VALENCE_STD[a] for a in atoms])\n",
    "        n_charge = np.zeros(len(atoms), dtype=np.float16)\n",
    "        isleaf   = np.zeros(len(atoms), dtype=bool) # is the atom in the leafs of connection tree?\n",
    "        coords   = molecule[['x', 'y', 'z']].values\n",
    "        kdt = KDTree(coords) # use an optimized structure for closest match query\n",
    "        nbond = {}\n",
    "        connected = { i:{} for i in atoms_idx}\n",
    "        \n",
    "        # select Hydrogen first to avoid butadyne-like ordering failures (molecule_name=dsgdb9nsd_000023)\n",
    "        ordered_atoms_index = list(atoms_idx)\n",
    "        ordered_atoms_index.sort(key = lambda i: BOND_ORDER[atoms[i]])\n",
    "        ordered_atoms_index = np.asarray(ordered_atoms_index)\n",
    "        \n",
    "        # STEP 1: 1-bond connect each atom with closest match\n",
    "        #         only one bond for each atom pair is done in step 1\n",
    "        for a0 in ordered_atoms_index:\n",
    "            search_bonds(kdt, n_avail, nbond, connected, isleaf, coords, atoms, atoms_idx, \n",
    "                         a0, connect_once=True, VALENCE=VALENCE_STD)\n",
    "     \n",
    "        # STEP 2: greedy connect n-bonds, progressing from leafs of connection tree\n",
    "        while (((n_avail > 0).sum() > 0) and isleaf).sum() > 0:\n",
    "            progress = False\n",
    "            for a0 in ordered_atoms_index:\n",
    "                #print(\"leaF/Trunk & avail: \" + \", \".join([f\"{i}:{atoms[i]}={leaflabel[leaf[i]]}{n_avail[i]}\"\n",
    "                #                                          for i in ordered_atoms_index]))\n",
    "                if (n_avail[a0] > 0) and isleaf[a0]:\n",
    "                    for a1 in connected[a0]:\n",
    "                        if (n_avail[a0] > 0) and (n_avail[a1] > 0):\n",
    "                            add_bond(n_avail, nbond, a0, a1) \n",
    "                            progress = True\n",
    "                            if (n_avail[a0] == 0) or (n_avail[a1] == 0):\n",
    "                                isleaf[a0] = 1\n",
    "                                isleaf[a1] = 1\n",
    "            if not progress:\n",
    "                break\n",
    "        \n",
    "        # gather remaining multiple bonds\n",
    "        if n_avail.sum() > 0:\n",
    "            for key in nbond.keys():\n",
    "                a0, a1 = key\n",
    "                while (n_avail[a0] > 0) and (n_avail[a1] > 0):\n",
    "                    add_bond(n_avail, nbond, a0, a1)\n",
    "                    \n",
    "        # STEP 3: Detect cycles : algorithm complexity in O(m^2 * n)\n",
    "        # paper : https://link.springer.com/article/10.1007/s00453-007-9064-z\n",
    "        # nx doc: https://networkx.github.io/documentation/latest/reference/algorithms/generated/networkx.algorithms.cycles.minimum_cycle_basis.html\n",
    "        graph = nx.Graph([bond for bond in nbond.keys()])\n",
    "        unordered_cycles = nx.minimum_cycle_basis(graph)\n",
    "        \n",
    "        # index atoms by their sequential order in the cycle: i.e follow bonds\n",
    "        # Note: this code can be written in a much cleaner way!\n",
    "        if len(unordered_cycles) > 0:\n",
    "            for icycle, c in enumerate(unordered_cycles):\n",
    "                available = {i:1 for i in c}\n",
    "                a0 = c[0]\n",
    "                cycle = [ a0 ]\n",
    "                del(available[a0])\n",
    "                for index in range(1, len(c)):\n",
    "                    # get atoms bonded to a0\n",
    "                    bonded = [ b for b in nbond.keys() if a0 in b ]\n",
    "                    bonded = list(map(lambda b: b[0] if b[1] == a0 else b[1], bonded))\n",
    "                    \n",
    "                    # get next atom and remove it from cycle\n",
    "                    assert(len(bonded) > 0)\n",
    "                    found = False\n",
    "                    for a1 in bonded:\n",
    "                        if (a1 in bonded) and (a1 in available):\n",
    "                            cycle.append(a1)\n",
    "                            del(available[a1])\n",
    "                            a0 = a1\n",
    "                            found = True\n",
    "                            break\n",
    "                    assert(found)\n",
    "                    \n",
    "                # and add cycles found to the cycle dataframe lists\n",
    "                cycle_name.extend([name] * len(cycle))\n",
    "                cycle_index.extend([icycle] * len(cycle))\n",
    "                cycle_seq.extend(np.arange(len(cycle)))\n",
    "                cycle_atom_index.extend(cycle)\n",
    "        \n",
    "        # display info on failed molecules\n",
    "        if n_avail.sum() > 0:\n",
    "            error = 1\n",
    "            print(f\"   Remaining bondings={n_avail.sum()} for molecule_name={name}, atoms: \" +\n",
    "                  \", \".join([f\"{i}:{atoms[i]}\" for i in atoms_idx if n_avail[i] > 0]))\n",
    "\n",
    "        # inputs for DataFrame bonds\n",
    "        for (a0, a1), (n, dist) in nbond.items():\n",
    "            # append to python lists which is 7x faster than toa pd.DataFrame\n",
    "            out_name.append(name)\n",
    "            out_a0.append(a0)\n",
    "            out_a1.append(a1)\n",
    "            out_n.append(n)\n",
    "            out_dist.append(dist)\n",
    "            out_error.append(error)\n",
    "            out_type.append(f\"{n:0.1f}\" + \"\".join(sorted(f\"{atoms[a0]}{atoms[a1]}\")) )\n",
    "    \n",
    "        # inputs for DataFrame charges\n",
    "        charge_name.extend([name]*len(atoms))\n",
    "        charge_atom_index.extend(molecule.atom_index.values)\n",
    "        charge_value.extend(n_charge)\n",
    "        \n",
    "    bonds = pd.DataFrame({'molecule_name':out_name, 'atom_index_0':out_a0,'atom_index_1':out_a1, 'nbond':out_n, \n",
    "                          'L2dist':out_dist, 'error':out_error, 'bond_type':out_type})\n",
    "    cycles = pd.DataFrame({'molecule_name' : cycle_name, 'cycle_index' : cycle_index, \n",
    "                           'cycle_seq' : cycle_seq, 'atom_index' : cycle_atom_index})\n",
    "    \n",
    "    return bonds, cycles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8c10a7d5-8806-49c7-ab2f-9ff87564964a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16e49a04b5294330893cbfdc9d3870de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000204, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000208, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000731, atoms: 2:C, 4:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000761, atoms: 2:C, 4:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000854, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000909, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000911, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000912, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000968, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000974, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000982, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001005, atoms: 0:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001006, atoms: 0:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001012, atoms: 0:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001013, atoms: 0:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001014, atoms: 0:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001046, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001047, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001050, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001052, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001053, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001076, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001642, atoms: 6:O, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001855, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_002869, atoms: 0:C, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_002872, atoms: 0:N, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003412, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003414, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003632, atoms: 0:C, 5:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003634, atoms: 0:C, 5:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003717, atoms: 3:C, 5:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003766, atoms: 2:N, 6:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003767, atoms: 2:N, 6:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003769, atoms: 2:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003770, atoms: 2:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003813, atoms: 0:N, 5:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003844, atoms: 3:N, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003845, atoms: 3:N, 6:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003894, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003904, atoms: 3:N, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003905, atoms: 3:N, 6:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003952, atoms: 3:C, 5:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004261, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004262, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004268, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004273, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004279, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004280, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004287, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004295, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004298, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004302, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004311, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004313, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004323, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004325, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004327, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004468, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004472, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004482, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004517, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004518, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004596, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004602, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004611, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004615, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004618, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004621, atoms: 5:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004673, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004675, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004720, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004724, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004728, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004732, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004736, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004739, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004741, atoms: 2:C, 6:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004743, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004747, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004749, atoms: 2:C, 6:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004767, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004768, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004770, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004788, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004790, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004797, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004799, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004800, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004801, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004814, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004849, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004850, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004856, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004862, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004863, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004864, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004865, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004912, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004914, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004942, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005050, atoms: 0:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005051, atoms: 0:O, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005052, atoms: 0:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005053, atoms: 0:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005055, atoms: 0:O, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005056, atoms: 0:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005146, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005148, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005159, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005164, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005165, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005177, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005183, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005186, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005196, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005198, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005207, atoms: 1:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005208, atoms: 1:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005217, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005228, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005229, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005235, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005252, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005264, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005305, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005306, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005315, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005316, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005324, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005332, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005334, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005452, atoms: 4:C, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005453, atoms: 4:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005456, atoms: 4:C, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005457, atoms: 4:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005461, atoms: 4:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005464, atoms: 4:C, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005470, atoms: 4:C, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005471, atoms: 4:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005475, atoms: 4:C, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005477, atoms: 4:C, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005479, atoms: 4:C, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005481, atoms: 4:C, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005519, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005554, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005569, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005572, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005574, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005610, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005611, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005650, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005654, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005658, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005678, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005679, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005695, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005697, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005704, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005705, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005706, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005710, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005712, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005734, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005736, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005765, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005766, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005780, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005782, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005788, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005792, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005794, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005804, atoms: 1:C, 10:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006033, atoms: 3:C, 10:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006059, atoms: 8:O, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006075, atoms: 1:C, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006170, atoms: 1:C, 10:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006180, atoms: 5:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006182, atoms: 5:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_006683, atoms: 2:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007326, atoms: 8:O, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007335, atoms: 6:O, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007654, atoms: 5:C, 8:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007657, atoms: 4:C, 10:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007699, atoms: 7:O, 9:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007907, atoms: 4:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007914, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007918, atoms: 4:O, 6:C\n"
     ]
    }
   ],
   "source": [
    "train_bonds, train_cycles = compute_bonds(structures.set_index('molecule_name'), train.molecule_name.unique()[:5000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f8cadf1a-ae78-4758-adad-37c624df411a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c2f7b1c3711459c8ca9833b42017928",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000202, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000206, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000210, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000211, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000271, atoms: 4:C, 8:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000282, atoms: 6:O, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000724, atoms: 2:N, 5:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000782, atoms: 2:C, 4:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000840, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000842, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000843, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000855, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_000910, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001007, atoms: 0:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001008, atoms: 0:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001011, atoms: 0:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001062, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001064, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001074, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001648, atoms: 7:O, 8:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001670, atoms: 2:C, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_001853, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003411, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003413, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003765, atoms: 2:N, 6:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003768, atoms: 2:N, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_003893, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004269, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004272, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004286, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004294, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004300, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004315, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004322, atoms: 0:O, 3:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004328, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004470, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004599, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004605, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004608, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004620, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004671, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004721, atoms: 2:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004725, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004729, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004733, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004737, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004738, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004745, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004750, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004789, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004791, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004798, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004802, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004804, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004806, atoms: 2:C, 4:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004844, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004848, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004861, atoms: 2:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004907, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004910, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_004915, atoms: 3:C, 5:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005054, atoms: 0:O, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005150, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005152, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005182, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005184, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005188, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005192, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005194, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005200, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005205, atoms: 1:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005206, atoms: 1:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005232, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005234, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005250, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005262, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005266, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005268, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005286, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005287, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005322, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005336, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005460, atoms: 4:C, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005465, atoms: 4:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005468, atoms: 4:C, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005469, atoms: 4:C, 7:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005473, atoms: 4:C, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005482, atoms: 4:C, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005517, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005536, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005556, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005567, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005601, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005651, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005655, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005661, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005662, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005703, atoms: 3:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005738, atoms: 2:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_005790, atoms: 4:C, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007375, atoms: 4:C, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007429, atoms: 8:O, 9:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007602, atoms: 3:C, 10:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007909, atoms: 4:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007913, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_007917, atoms: 4:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_010016, atoms: 1:C, 8:N\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_011382, atoms: 5:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_011384, atoms: 5:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_011390, atoms: 5:O, 7:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_012586, atoms: 3:O, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_012589, atoms: 3:O, 7:O\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_013833, atoms: 2:O, 9:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_014316, atoms: 0:O, 6:C\n",
      "   Remaining bondings=2 for molecule_name=dsgdb9nsd_014318, atoms: 0:O, 6:C\n"
     ]
    }
   ],
   "source": [
    "test_bonds, test_cycles = compute_bonds(structures.set_index('molecule_name'), test.molecule_name.unique()[:5000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6db33058-af07-4110-b6a7-2282a29b87a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 3., 2.])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.unique(train_bonds.nbond)\n",
    "pd.unique(test_bonds.nbond)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f9c81dee-8cce-4adb-bf02-ce10e59c7f0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_bonds.to_csv('train_bonds.csv', index=False)\n",
    "train_cycles.to_csv('train_cycles.csv', index=False)\n",
    "\n",
    "test_bonds.to_csv('test_bonds.csv', index=False)\n",
    "test_cycles.to_csv('test_cycles.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('py38')",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "f2b0cf5112dfd2af73b780f4815bf58f6493434a6fee69d3bee328e26a09a438"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}