gen_input_graph_qm9.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "likely-profession",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.spatial import distance_matrix\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "married-tobago",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>mol_id</th>\n",
       "      <th>A</th>\n",
       "      <th>B</th>\n",
       "      <th>C</th>\n",
       "      <th>mu</th>\n",
       "      <th>alpha</th>\n",
       "      <th>homo</th>\n",
       "      <th>lumo</th>\n",
       "      <th>gap</th>\n",
       "      <th>r2</th>\n",
       "      <th>zpve</th>\n",
       "      <th>u0</th>\n",
       "      <th>u298</th>\n",
       "      <th>h298</th>\n",
       "      <th>g298</th>\n",
       "      <th>cv</th>\n",
       "      <th>u0_atom</th>\n",
       "      <th>u298_atom</th>\n",
       "      <th>h298_atom</th>\n",
       "      <th>g298_atom</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>gdb_1</td>\n",
       "      <td>157.71180</td>\n",
       "      <td>157.709970</td>\n",
       "      <td>157.706990</td>\n",
       "      <td>0.0000</td>\n",
       "      <td>13.21</td>\n",
       "      <td>-0.3877</td>\n",
       "      <td>0.1171</td>\n",
       "      <td>0.5048</td>\n",
       "      <td>35.3641</td>\n",
       "      <td>0.044749</td>\n",
       "      <td>-40.478930</td>\n",
       "      <td>-40.476062</td>\n",
       "      <td>-40.475117</td>\n",
       "      <td>-40.498597</td>\n",
       "      <td>6.469</td>\n",
       "      <td>-395.999595</td>\n",
       "      <td>-398.643290</td>\n",
       "      <td>-401.014647</td>\n",
       "      <td>-372.471772</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>gdb_2</td>\n",
       "      <td>293.60975</td>\n",
       "      <td>293.541110</td>\n",
       "      <td>191.393970</td>\n",
       "      <td>1.6256</td>\n",
       "      <td>9.46</td>\n",
       "      <td>-0.2570</td>\n",
       "      <td>0.0829</td>\n",
       "      <td>0.3399</td>\n",
       "      <td>26.1563</td>\n",
       "      <td>0.034358</td>\n",
       "      <td>-56.525887</td>\n",
       "      <td>-56.523026</td>\n",
       "      <td>-56.522082</td>\n",
       "      <td>-56.544961</td>\n",
       "      <td>6.316</td>\n",
       "      <td>-276.861363</td>\n",
       "      <td>-278.620271</td>\n",
       "      <td>-280.399259</td>\n",
       "      <td>-259.338802</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>gdb_3</td>\n",
       "      <td>799.58812</td>\n",
       "      <td>437.903860</td>\n",
       "      <td>282.945450</td>\n",
       "      <td>1.8511</td>\n",
       "      <td>6.31</td>\n",
       "      <td>-0.2928</td>\n",
       "      <td>0.0687</td>\n",
       "      <td>0.3615</td>\n",
       "      <td>19.0002</td>\n",
       "      <td>0.021375</td>\n",
       "      <td>-76.404702</td>\n",
       "      <td>-76.401867</td>\n",
       "      <td>-76.400922</td>\n",
       "      <td>-76.422349</td>\n",
       "      <td>6.002</td>\n",
       "      <td>-213.087624</td>\n",
       "      <td>-213.974294</td>\n",
       "      <td>-215.159658</td>\n",
       "      <td>-201.407171</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>gdb_4</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>35.610036</td>\n",
       "      <td>35.610036</td>\n",
       "      <td>0.0000</td>\n",
       "      <td>16.28</td>\n",
       "      <td>-0.2845</td>\n",
       "      <td>0.0506</td>\n",
       "      <td>0.3351</td>\n",
       "      <td>59.5248</td>\n",
       "      <td>0.026841</td>\n",
       "      <td>-77.308427</td>\n",
       "      <td>-77.305527</td>\n",
       "      <td>-77.304583</td>\n",
       "      <td>-77.327429</td>\n",
       "      <td>8.574</td>\n",
       "      <td>-385.501997</td>\n",
       "      <td>-387.237686</td>\n",
       "      <td>-389.016047</td>\n",
       "      <td>-365.800724</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>gdb_5</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>44.593883</td>\n",
       "      <td>44.593883</td>\n",
       "      <td>2.8937</td>\n",
       "      <td>12.99</td>\n",
       "      <td>-0.3604</td>\n",
       "      <td>0.0191</td>\n",
       "      <td>0.3796</td>\n",
       "      <td>48.7476</td>\n",
       "      <td>0.016601</td>\n",
       "      <td>-93.411888</td>\n",
       "      <td>-93.409370</td>\n",
       "      <td>-93.408425</td>\n",
       "      <td>-93.431246</td>\n",
       "      <td>6.278</td>\n",
       "      <td>-301.820534</td>\n",
       "      <td>-302.906752</td>\n",
       "      <td>-304.091489</td>\n",
       "      <td>-288.720028</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>133880</th>\n",
       "      <td>gdb_133881</td>\n",
       "      <td>3.59483</td>\n",
       "      <td>2.198990</td>\n",
       "      <td>1.904230</td>\n",
       "      <td>1.6637</td>\n",
       "      <td>69.37</td>\n",
       "      <td>-0.2254</td>\n",
       "      <td>0.0588</td>\n",
       "      <td>0.2842</td>\n",
       "      <td>760.7472</td>\n",
       "      <td>0.127406</td>\n",
       "      <td>-400.633868</td>\n",
       "      <td>-400.628599</td>\n",
       "      <td>-400.627654</td>\n",
       "      <td>-400.663098</td>\n",
       "      <td>23.658</td>\n",
       "      <td>-1603.983913</td>\n",
       "      <td>-1614.898804</td>\n",
       "      <td>-1623.788097</td>\n",
       "      <td>-1492.819438</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>133881</th>\n",
       "      <td>gdb_133882</td>\n",
       "      <td>3.65648</td>\n",
       "      <td>2.142370</td>\n",
       "      <td>1.904390</td>\n",
       "      <td>1.2976</td>\n",
       "      <td>69.52</td>\n",
       "      <td>-0.2393</td>\n",
       "      <td>0.0608</td>\n",
       "      <td>0.3002</td>\n",
       "      <td>762.6354</td>\n",
       "      <td>0.127495</td>\n",
       "      <td>-400.629713</td>\n",
       "      <td>-400.624444</td>\n",
       "      <td>-400.623500</td>\n",
       "      <td>-400.658942</td>\n",
       "      <td>23.697</td>\n",
       "      <td>-1601.376613</td>\n",
       "      <td>-1612.291504</td>\n",
       "      <td>-1621.181424</td>\n",
       "      <td>-1490.211511</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>133882</th>\n",
       "      <td>gdb_133883</td>\n",
       "      <td>3.67118</td>\n",
       "      <td>2.143140</td>\n",
       "      <td>1.895010</td>\n",
       "      <td>1.2480</td>\n",
       "      <td>73.60</td>\n",
       "      <td>-0.2233</td>\n",
       "      <td>0.0720</td>\n",
       "      <td>0.2953</td>\n",
       "      <td>780.3553</td>\n",
       "      <td>0.140458</td>\n",
       "      <td>-380.753918</td>\n",
       "      <td>-380.748619</td>\n",
       "      <td>-380.747675</td>\n",
       "      <td>-380.783148</td>\n",
       "      <td>23.972</td>\n",
       "      <td>-1667.045429</td>\n",
       "      <td>-1678.830048</td>\n",
       "      <td>-1688.312964</td>\n",
       "      <td>-1549.143391</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>133883</th>\n",
       "      <td>gdb_133884</td>\n",
       "      <td>3.52845</td>\n",
       "      <td>2.151310</td>\n",
       "      <td>1.865820</td>\n",
       "      <td>1.9576</td>\n",
       "      <td>77.40</td>\n",
       "      <td>-0.2122</td>\n",
       "      <td>0.0881</td>\n",
       "      <td>0.3003</td>\n",
       "      <td>803.1904</td>\n",
       "      <td>0.152222</td>\n",
       "      <td>-364.720374</td>\n",
       "      <td>-364.714974</td>\n",
       "      <td>-364.714030</td>\n",
       "      <td>-364.749650</td>\n",
       "      <td>24.796</td>\n",
       "      <td>-1794.600439</td>\n",
       "      <td>-1807.210860</td>\n",
       "      <td>-1817.286772</td>\n",
       "      <td>-1670.349892</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>133884</th>\n",
       "      <td>gdb_133885</td>\n",
       "      <td>3.64015</td>\n",
       "      <td>2.217640</td>\n",
       "      <td>1.937930</td>\n",
       "      <td>0.8626</td>\n",
       "      <td>69.48</td>\n",
       "      <td>-0.2316</td>\n",
       "      <td>0.0742</td>\n",
       "      <td>0.3058</td>\n",
       "      <td>756.3557</td>\n",
       "      <td>0.127862</td>\n",
       "      <td>-400.633052</td>\n",
       "      <td>-400.627892</td>\n",
       "      <td>-400.626948</td>\n",
       "      <td>-400.662186</td>\n",
       "      <td>23.434</td>\n",
       "      <td>-1603.471865</td>\n",
       "      <td>-1614.455155</td>\n",
       "      <td>-1623.345075</td>\n",
       "      <td>-1492.247150</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>133885 rows × 20 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            mol_id          A           B           C      mu  alpha    homo  \\\n",
       "0            gdb_1  157.71180  157.709970  157.706990  0.0000  13.21 -0.3877   \n",
       "1            gdb_2  293.60975  293.541110  191.393970  1.6256   9.46 -0.2570   \n",
       "2            gdb_3  799.58812  437.903860  282.945450  1.8511   6.31 -0.2928   \n",
       "3            gdb_4    0.00000   35.610036   35.610036  0.0000  16.28 -0.2845   \n",
       "4            gdb_5    0.00000   44.593883   44.593883  2.8937  12.99 -0.3604   \n",
       "...            ...        ...         ...         ...     ...    ...     ...   \n",
       "133880  gdb_133881    3.59483    2.198990    1.904230  1.6637  69.37 -0.2254   \n",
       "133881  gdb_133882    3.65648    2.142370    1.904390  1.2976  69.52 -0.2393   \n",
       "133882  gdb_133883    3.67118    2.143140    1.895010  1.2480  73.60 -0.2233   \n",
       "133883  gdb_133884    3.52845    2.151310    1.865820  1.9576  77.40 -0.2122   \n",
       "133884  gdb_133885    3.64015    2.217640    1.937930  0.8626  69.48 -0.2316   \n",
       "\n",
       "          lumo     gap        r2      zpve          u0        u298  \\\n",
       "0       0.1171  0.5048   35.3641  0.044749  -40.478930  -40.476062   \n",
       "1       0.0829  0.3399   26.1563  0.034358  -56.525887  -56.523026   \n",
       "2       0.0687  0.3615   19.0002  0.021375  -76.404702  -76.401867   \n",
       "3       0.0506  0.3351   59.5248  0.026841  -77.308427  -77.305527   \n",
       "4       0.0191  0.3796   48.7476  0.016601  -93.411888  -93.409370   \n",
       "...        ...     ...       ...       ...         ...         ...   \n",
       "133880  0.0588  0.2842  760.7472  0.127406 -400.633868 -400.628599   \n",
       "133881  0.0608  0.3002  762.6354  0.127495 -400.629713 -400.624444   \n",
       "133882  0.0720  0.2953  780.3553  0.140458 -380.753918 -380.748619   \n",
       "133883  0.0881  0.3003  803.1904  0.152222 -364.720374 -364.714974   \n",
       "133884  0.0742  0.3058  756.3557  0.127862 -400.633052 -400.627892   \n",
       "\n",
       "              h298        g298      cv      u0_atom    u298_atom    h298_atom  \\\n",
       "0       -40.475117  -40.498597   6.469  -395.999595  -398.643290  -401.014647   \n",
       "1       -56.522082  -56.544961   6.316  -276.861363  -278.620271  -280.399259   \n",
       "2       -76.400922  -76.422349   6.002  -213.087624  -213.974294  -215.159658   \n",
       "3       -77.304583  -77.327429   8.574  -385.501997  -387.237686  -389.016047   \n",
       "4       -93.408425  -93.431246   6.278  -301.820534  -302.906752  -304.091489   \n",
       "...            ...         ...     ...          ...          ...          ...   \n",
       "133880 -400.627654 -400.663098  23.658 -1603.983913 -1614.898804 -1623.788097   \n",
       "133881 -400.623500 -400.658942  23.697 -1601.376613 -1612.291504 -1621.181424   \n",
       "133882 -380.747675 -380.783148  23.972 -1667.045429 -1678.830048 -1688.312964   \n",
       "133883 -364.714030 -364.749650  24.796 -1794.600439 -1807.210860 -1817.286772   \n",
       "133884 -400.626948 -400.662186  23.434 -1603.471865 -1614.455155 -1623.345075   \n",
       "\n",
       "          g298_atom  \n",
       "0       -372.471772  \n",
       "1       -259.338802  \n",
       "2       -201.407171  \n",
       "3       -365.800724  \n",
       "4       -288.720028  \n",
       "...             ...  \n",
       "133880 -1492.819438  \n",
       "133881 -1490.211511  \n",
       "133882 -1549.143391  \n",
       "133883 -1670.349892  \n",
       "133884 -1492.247150  \n",
       "\n",
       "[133885 rows x 20 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = \"../qm9/\"\n",
    "qm9 = pd.read_csv(dataset + \"gdb9.sdf.csv\")\n",
    "qm9"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compatible-reset",
   "metadata": {},
   "source": [
    "## Find representations for nodes and edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dress-seminar",
   "metadata": {},
   "outputs": [],
   "source": [
    "def qm9_nodes(g, hydrogen=False):\n",
    "    h = []\n",
    "    for n, d in g.nodes_iter(data=True):\n",
    "        h_t = []\n",
    "        # Atom type (One-hot H, C, N, O F)\n",
    "        h_t += [int(d['a_type'] == x) for x in ['H', 'C', 'N', 'O', 'F']]\n",
    "        # Atomic number\n",
    "        h_t.append(d['a_num'])\n",
    "        # Partial Charge\n",
    "        h_t.append(d['pc'])\n",
    "        # Acceptor\n",
    "        h_t.append(d['acceptor'])\n",
    "        # Donor\n",
    "        h_t.append(d['donor'])\n",
    "        # Aromatic\n",
    "        h_t.append(int(d['aromatic']))\n",
    "        # Hybradization\n",
    "        h_t += [int(d['hybridization'] == x) for x in [rdkit.Chem.rdchem.HybridizationType.SP, rdkit.Chem.rdchem.HybridizationType.SP2, rdkit.Chem.rdchem.HybridizationType.SP3]]\n",
    "        # If number hydrogen is used as a\n",
    "        if hydrogen:\n",
    "            h_t.append(d['num_h'])\n",
    "        h.append(h_t)\n",
    "    return h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fluid-space",
   "metadata": {},
   "source": [
    "## Normalize scalar coupling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "happy-admission",
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_min  = train['scalar_coupling_constant'].min()\n",
    "scale_max  = train['scalar_coupling_constant'].max()\n",
    "scale_mid = (scale_max + scale_min)/2\n",
    "scale_norm = scale_max - scale_mid\n",
    "\n",
    "train['scalar_coupling_constant'] = (train['scalar_coupling_constant'] - scale_mid)/scale_norm\n",
    "\n",
    "train[['1JHC', '1JHN', '2JHC', '2JHH', '2JHN', '3JHC', '3JHH', '3JHN']] = pd.get_dummies(train['type'])\n",
    "test[['1JHC', '1JHN', '2JHC', '2JHH', '2JHN', '3JHC', '3JHH', '3JHN']] = pd.get_dummies(test['type'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "applied-going",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>molecule_name</th>\n",
       "      <th>atom_index_0</th>\n",
       "      <th>atom_index_1</th>\n",
       "      <th>type</th>\n",
       "      <th>scalar_coupling_constant</th>\n",
       "      <th>1JHC</th>\n",
       "      <th>1JHN</th>\n",
       "      <th>2JHC</th>\n",
       "      <th>2JHH</th>\n",
       "      <th>2JHN</th>\n",
       "      <th>3JHC</th>\n",
       "      <th>3JHH</th>\n",
       "      <th>3JHN</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1JHC</td>\n",
       "      <td>0.026406</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2JHH</td>\n",
       "      <td>-0.734594</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>2JHH</td>\n",
       "      <td>-0.734576</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>2JHH</td>\n",
       "      <td>-0.734572</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>1JHC</td>\n",
       "      <td>0.026404</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4659071</th>\n",
       "      <td>4659071</td>\n",
       "      <td>dsgdb9nsd_133884</td>\n",
       "      <td>17</td>\n",
       "      <td>4</td>\n",
       "      <td>2JHC</td>\n",
       "      <td>-0.617348</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4659072</th>\n",
       "      <td>4659072</td>\n",
       "      <td>dsgdb9nsd_133884</td>\n",
       "      <td>17</td>\n",
       "      <td>5</td>\n",
       "      <td>3JHC</td>\n",
       "      <td>-0.640911</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4659073</th>\n",
       "      <td>4659073</td>\n",
       "      <td>dsgdb9nsd_133884</td>\n",
       "      <td>17</td>\n",
       "      <td>6</td>\n",
       "      <td>3JHC</td>\n",
       "      <td>-0.636123</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4659074</th>\n",
       "      <td>4659074</td>\n",
       "      <td>dsgdb9nsd_133884</td>\n",
       "      <td>17</td>\n",
       "      <td>7</td>\n",
       "      <td>2JHC</td>\n",
       "      <td>-0.607695</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4659075</th>\n",
       "      <td>4659075</td>\n",
       "      <td>dsgdb9nsd_133884</td>\n",
       "      <td>17</td>\n",
       "      <td>8</td>\n",
       "      <td>1JHC</td>\n",
       "      <td>0.288825</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4659076 rows × 14 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "              id     molecule_name  atom_index_0  atom_index_1  type  \\\n",
       "0              0  dsgdb9nsd_000001             1             0  1JHC   \n",
       "1              1  dsgdb9nsd_000001             1             2  2JHH   \n",
       "2              2  dsgdb9nsd_000001             1             3  2JHH   \n",
       "3              3  dsgdb9nsd_000001             1             4  2JHH   \n",
       "4              4  dsgdb9nsd_000001             2             0  1JHC   \n",
       "...          ...               ...           ...           ...   ...   \n",
       "4659071  4659071  dsgdb9nsd_133884            17             4  2JHC   \n",
       "4659072  4659072  dsgdb9nsd_133884            17             5  3JHC   \n",
       "4659073  4659073  dsgdb9nsd_133884            17             6  3JHC   \n",
       "4659074  4659074  dsgdb9nsd_133884            17             7  2JHC   \n",
       "4659075  4659075  dsgdb9nsd_133884            17             8  1JHC   \n",
       "\n",
       "         scalar_coupling_constant  1JHC  1JHN  2JHC  2JHH  2JHN  3JHC  3JHH  \\\n",
       "0                        0.026406     1     0     0     0     0     0     0   \n",
       "1                       -0.734594     0     0     0     1     0     0     0   \n",
       "2                       -0.734576     0     0     0     1     0     0     0   \n",
       "3                       -0.734572     0     0     0     1     0     0     0   \n",
       "4                        0.026404     1     0     0     0     0     0     0   \n",
       "...                           ...   ...   ...   ...   ...   ...   ...   ...   \n",
       "4659071                 -0.617348     0     0     1     0     0     0     0   \n",
       "4659072                 -0.640911     0     0     0     0     0     1     0   \n",
       "4659073                 -0.636123     0     0     0     0     0     1     0   \n",
       "4659074                 -0.607695     0     0     1     0     0     0     0   \n",
       "4659075                  0.288825     1     0     0     0     0     0     0   \n",
       "\n",
       "         3JHN  \n",
       "0           0  \n",
       "1           0  \n",
       "2           0  \n",
       "3           0  \n",
       "4           0  \n",
       "...       ...  \n",
       "4659071     0  \n",
       "4659072     0  \n",
       "4659073     0  \n",
       "4659074     0  \n",
       "4659075     0  \n",
       "\n",
       "[4659076 rows x 14 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ordered-bolivia",
   "metadata": {},
   "source": [
    "## One hot encoding atom type and normalize distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bulgarian-montgomery",
   "metadata": {},
   "outputs": [],
   "source": [
    "structures[['C', 'F' ,'H', 'N', 'O']] = pd.get_dummies(structures['atom'])\n",
    "structures[['x', 'y', 'z']] = structures[['x', 'y', 'z']]/10."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mineral-seven",
   "metadata": {},
   "source": [
    "## Process bonds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "brazilian-printing",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>molecule_name</th>\n",
       "      <th>atom_index_0</th>\n",
       "      <th>atom_index_1</th>\n",
       "      <th>nbond</th>\n",
       "      <th>L2dist</th>\n",
       "      <th>error</th>\n",
       "      <th>bond_type</th>\n",
       "      <th>nbond_1</th>\n",
       "      <th>nbond_2</th>\n",
       "      <th>nbond_3</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.091953</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CH</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.091952</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CH</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.091946</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CH</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>dsgdb9nsd_000001</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.091948</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CH</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>dsgdb9nsd_000002</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.017190</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0HN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76653</th>\n",
       "      <td>dsgdb9nsd_007969</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.542051</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CC</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76654</th>\n",
       "      <td>dsgdb9nsd_007969</td>\n",
       "      <td>1</td>\n",
       "      <td>6</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.549548</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CC</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76655</th>\n",
       "      <td>dsgdb9nsd_007969</td>\n",
       "      <td>2</td>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.513276</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CC</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76656</th>\n",
       "      <td>dsgdb9nsd_007969</td>\n",
       "      <td>3</td>\n",
       "      <td>5</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.357137</td>\n",
       "      <td>0</td>\n",
       "      <td>2.0CC</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76657</th>\n",
       "      <td>dsgdb9nsd_007969</td>\n",
       "      <td>5</td>\n",
       "      <td>6</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.454146</td>\n",
       "      <td>0</td>\n",
       "      <td>1.0CC</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>76658 rows × 10 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          molecule_name  atom_index_0  atom_index_1  nbond    L2dist  error  \\\n",
       "0      dsgdb9nsd_000001             0             1    1.0  1.091953      0   \n",
       "1      dsgdb9nsd_000001             0             2    1.0  1.091952      0   \n",
       "2      dsgdb9nsd_000001             0             3    1.0  1.091946      0   \n",
       "3      dsgdb9nsd_000001             0             4    1.0  1.091948      0   \n",
       "4      dsgdb9nsd_000002             0             1    1.0  1.017190      0   \n",
       "...                 ...           ...           ...    ...       ...    ...   \n",
       "76653  dsgdb9nsd_007969             1             2    1.0  1.542051      0   \n",
       "76654  dsgdb9nsd_007969             1             6    1.0  1.549548      0   \n",
       "76655  dsgdb9nsd_007969             2             3    1.0  1.513276      0   \n",
       "76656  dsgdb9nsd_007969             3             5    2.0  1.357137      0   \n",
       "76657  dsgdb9nsd_007969             5             6    1.0  1.454146      0   \n",
       "\n",
       "      bond_type  nbond_1  nbond_2  nbond_3  \n",
       "0         1.0CH        1        0        0  \n",
       "1         1.0CH        1        0        0  \n",
       "2         1.0CH        1        0        0  \n",
       "3         1.0CH        1        0        0  \n",
       "4         1.0HN        1        0        0  \n",
       "...         ...      ...      ...      ...  \n",
       "76653     1.0CC        1        0        0  \n",
       "76654     1.0CC        1        0        0  \n",
       "76655     1.0CC        1        0        0  \n",
       "76656     2.0CC        0        1        0  \n",
       "76657     1.0CC        1        0        0  \n",
       "\n",
       "[76658 rows x 10 columns]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_bonds[['nbond_1', 'nbond_2', 'nbond_3']] = pd.get_dummies(train_bonds['nbond'])#train_bonds['nbond']/3\n",
    "test_bonds[['nbond_1', 'nbond_2', 'nbond_3']] = pd.get_dummies(test_bonds['nbond'])#test_bonds['nbond']/3\n",
    "train_bonds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mysterious-luxembourg",
   "metadata": {},
   "source": [
    "## Process angle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "congressional-variation",
   "metadata": {},
   "outputs": [],
   "source": [
    "angles['dihedral'] = angles['dihedral']/np.pi\n",
    "angles['shortest_path_n_bonds'] = angles['shortest_path_n_bonds']/6.0\n",
    "angles = angles.fillna(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acute-pressure",
   "metadata": {},
   "source": [
    "## Split structures to train and test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "secure-greensboro",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "29\n"
     ]
    }
   ],
   "source": [
    "train_mol_names = train['molecule_name'].unique()\n",
    "test_mol_names  = test['molecule_name'].unique()\n",
    "\n",
    "train_structures = structures.loc[structures['molecule_name'].isin(train_mol_names)]\n",
    "test_structures = structures.loc[structures['molecule_name'].isin(test_mol_names)]\n",
    "\n",
    "train_struct_group = train_structures.groupby('molecule_name')\n",
    "test_struct_group  = test_structures.groupby('molecule_name')\n",
    "\n",
    "train_group = train.groupby('molecule_name')\n",
    "test_group  = test.groupby('molecule_name')\n",
    "\n",
    "train_bond_group = train_bonds.groupby('molecule_name')\n",
    "test_bond_group  = test_bonds.groupby('molecule_name')\n",
    "\n",
    "train_angs = angles.loc[angles['molecule_name'].isin(train_mol_names)]\n",
    "test_angs = angles.loc[angles['molecule_name'].isin(test_mol_names)]\n",
    "\n",
    "train_angs_group = train_angs.groupby('molecule_name')\n",
    "test_angs_group  = test_angs.groupby('molecule_name')\n",
    "\n",
    "# Find max nodes in graph:\n",
    "max_size = train_struct_group.size().max()\n",
    "print(max_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interim-angola",
   "metadata": {},
   "source": [
    "## Define node and edge values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "positive-evolution",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Values for nodes\n",
    "node_vals = ['C', 'F' ,'H', 'N', 'O']#, 'x', 'y', 'z']\n",
    "# Values for edges\n",
    "bond_vals = ['nbond_1', 'nbond_2', 'nbond_3']\n",
    "ang_vals = ['shortest_path_n_bonds', 'cosinus', 'dihedral']\n",
    "j_coup_vals = ['1JHC', '1JHN', '2JHC', '2JHH', '2JHN', '3JHC', '3JHH', '3JHN']\n",
    "edge_vals = j_coup_vals + bond_vals + ang_vals\n",
    "\n",
    "# Find amount of training molecules\n",
    "n_train_mols = len(train_mol_names)\n",
    "n_test_mols = len(test_mol_names)\n",
    "\n",
    "# Find dimension of nodes and edges\n",
    "node_dim  = len(node_vals)\n",
    "bond_dim  = len(bond_vals)\n",
    "ang_dim   = len(ang_vals)\n",
    "j_coup_dim= len(j_coup_vals)\n",
    "\n",
    "edge_dim  = len(edge_vals) \n",
    "\n",
    "# Additional edge dims for distances \n",
    "add_edge_dim = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "shaped-directive",
   "metadata": {},
   "source": [
    "## Pre-allocate array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "affecting-mexico",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_nodes_array = np.zeros((n_train_mols, max_size, node_dim), dtype=np.float32) \n",
    "train_in_edges_array = np.zeros((n_train_mols, max_size, max_size, edge_dim + add_edge_dim), dtype=np.float32) \n",
    "train_out_edges_array = np.zeros((n_train_mols, max_size, max_size, 1), dtype=np.float32) \n",
    "\n",
    "test_nodes_array = np.zeros((n_test_mols, max_size, node_dim), dtype=np.float32) \n",
    "test_in_edges_array = np.zeros((n_test_mols, max_size, max_size, edge_dim + add_edge_dim), dtype=np.float32) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "motivated-delaware",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_arrs(val_group, struct_group, bond_group, ang_group, mode):\n",
    "    for i, (values, structs, bonds, angles) in tqdm(enumerate(zip(val_group, struct_group, bond_group, ang_group))):\n",
    "        \n",
    "        distances = np.zeros((max_size, max_size, add_edge_dim))\n",
    "        coords = structs[1][['x','y','z']].values\n",
    "        \n",
    "        dists  = distance_matrix(coords, coords)\n",
    "        \n",
    "        distances[:dists.shape[0],:dists.shape[1], 0] = dists\n",
    "        \n",
    "        # Create nodes\n",
    "        mol_info = structs[1][node_vals].values\n",
    "        nodes = np.zeros((max_size, node_dim))\n",
    "        nodes[:mol_info.shape[0], :mol_info.shape[1]] = mol_info\n",
    "\n",
    "        # Create edges\n",
    "        # in_feats is type descriptos one_hot_encoded -> use it to filter on type later on\n",
    "        in_feats = np.zeros((max_size, max_size, j_coup_dim))\n",
    "        ind = values[1][['atom_index_0', 'atom_index_1' ]].values\n",
    "        in_feats[ind[:,0], ind[:,1], 0:j_coup_dim] = values[1][j_coup_vals].values\n",
    "        in_feats[ind[:,1], ind[:,0], 0:j_coup_dim] = in_feats[ind[:,0], ind[:,1], 0:j_coup_dim]\n",
    "                  \n",
    "        # Create bonds\n",
    "        in_bonds = np.zeros((max_size, max_size, bond_dim))\n",
    "        ind_bonds = bonds[1][['atom_index_0', 'atom_index_1' ]].values\n",
    "        in_bonds[ind_bonds[:,0], ind_bonds[:,1]] = bonds[1][bond_vals].values\n",
    "        in_bonds[ind_bonds[:,1], ind_bonds[:,0]] = in_bonds[ind_bonds[:,0], ind_bonds[:,1]]\n",
    "        \n",
    "        # Create angles\n",
    "        ind_angs = angles[1][['atom_index_0', 'atom_index_1' ]].values\n",
    "        ang_mat  = np.zeros((max_size, max_size, ang_dim))\n",
    "        ang_mat[ind_angs[:,0], ind_angs[:,1]]  = angles[1][ang_vals]\n",
    "        ang_mat[ind_angs[:,1], ind_angs[:,0]]  = ang_mat[ind_angs[:,0], ind_angs[:,1]]\n",
    "        \n",
    "        # concat all edge values\n",
    "        in_edges = np.concatenate((in_feats, in_bonds, ang_mat, distances),axis=2)\n",
    "        \n",
    "        # create list of input array\n",
    "        if mode == \"train\":\n",
    "            out_edges = np.zeros((max_size, max_size, 1))\n",
    "            out_edges[ind[:,0], ind[:,1], 0] = values[1]['scalar_coupling_constant'].values\n",
    "            out_edges[ind[:,1], ind[:,0], 0] = out_edges[ind[:,0], ind[:,1], 0]\n",
    "        \n",
    "            train_nodes_array[i]      = nodes\n",
    "            train_in_edges_array[i]   = in_edges\n",
    "            train_out_edges_array[i]  = out_edges\n",
    "        elif mode == \"test\":\n",
    "            test_nodes_array[i]      = nodes\n",
    "            test_in_edges_array[i]   = in_edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "restricted-attack",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "311it [00:01, 223.02it/s]\n"
     ]
    }
   ],
   "source": [
    "make_arrs(train_group, train_struct_group, train_bond_group, train_angs_group, mode=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "renewable-asian",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "311it [00:01, 217.03it/s]\n"
     ]
    }
   ],
   "source": [
    "make_arrs(train_group, train_struct_group, train_bond_group, train_angs_group, mode=\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "frank-sweet",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(85012, 29, 5)\n",
      "(85012, 29, 29, 15)\n",
      "(85012, 29, 29, 1)\n",
      "(45777, 29, 5)\n",
      "(45777, 29, 29, 15)\n"
     ]
    }
   ],
   "source": [
    "print(train_nodes_array.shape)\n",
    "print(train_in_edges_array.shape)\n",
    "print(train_out_edges_array.shape)\n",
    "print(test_nodes_array.shape)\n",
    "print(test_in_edges_array.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "amateur-disclaimer",
   "metadata": {},
   "source": [
    "## Save arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "decimal-pipeline",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez_compressed(\"nodes_train.npz\" , train_nodes_array)\n",
    "np.savez_compressed(\"in_edges_train.npz\" , train_in_edges_array)\n",
    "np.savez_compressed(\"out_edges_train.npz\" , train_out_edges_array)\n",
    "\n",
    "np.savez_compressed(\"nodes_test.npz\" , test_nodes_array)\n",
    "np.savez_compressed(\"in_edges_test.npz\" , test_in_edges_array)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "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.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}