dGCN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "dGCN.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "CEUFzE4j8WEq",
        "34HBVTh68jHR",
        "Dd3rtzqR-hTy",
        "xuu4sTDFW4v3",
        "LIQBZjmdZnay",
        "PwgqP7Gxv2HS"
      ],
      "mount_file_id": "1AERyupaWbPJ2fOAJMSBR6w2ZAqXQozn3",
      "authorship_tag": "ABX9TyPq9SbtECBmUYL+cP1x+Etj",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/carlosmatherson/PolymerGCN/blob/main/dGCN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n"
      ],
      "metadata": {
        "id": "CEUFzE4j8WEq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install Spektral and RDKit."
      ],
      "metadata": {
        "id": "CSYHeRogUOQK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install spektral\n",
        "!pip install rdkit"
      ],
      "metadata": {
        "id": "mKWD-j7lBFaU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import necessary packages and functions."
      ],
      "metadata": {
        "id": "8CbOumsVUTNT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# general tools\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "# RDkit\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem.rdmolops import GetAdjacencyMatrix\n",
        "\n",
        "# Spektral\n",
        "from spektral.utils.sparse import reorder\n",
        "from spektral.data.graph import Graph\n",
        "from spektral.data import Dataset, BatchLoader\n",
        "from spektral.layers import GCNConv, GlobalSumPool\n",
        "\n",
        "# SciKit-Learn\n",
        "from sklearn.model_selection import KFold\n",
        "from sklearn.decomposition import PCA\n",
        "\n",
        "# TensorFlow & Keras\n",
        "from tensorflow.keras.layers import Dense, BatchNormalization\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.metrics import RootMeanSquaredError\n",
        "from tensorflow.keras import regularizers\n",
        "from tensorflow.keras import backend as K\n",
        "\n",
        "# Matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.colors import Normalize"
      ],
      "metadata": {
        "id": "MXbBM1BTqoyB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Atom Featurization"
      ],
      "metadata": {
        "id": "34HBVTh68jHR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Onehot Encoding: Maps input elements x which are not in the permitted list to the last element of the permitted list."
      ],
      "metadata": {
        "id": "xFd7WAE59HM4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Onehot Encoding\n",
        "def one_hot_encoding(x, permitted_list):\n",
        "\n",
        "    if x not in permitted_list:\n",
        "        x = permitted_list[-1]\n",
        "\n",
        "    binary_encoding = [int(boolean_value) for boolean_value in list(map(lambda s: x == s, permitted_list))]\n",
        "\n",
        "    return binary_encoding"
      ],
      "metadata": {
        "id": "2wkSe_gZ8DVD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Atom Featurization: Takes an RDKit atom object as input and gives a 1d-numpy array of atom features as output."
      ],
      "metadata": {
        "id": "EaWZf3teyKBP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Atom Featurization\n",
        "def get_atom_features(atom):\n",
        "\n",
        "    # Define list of permitted atoms\n",
        "    permitted_list_of_atoms = ['Br', 'C', 'Cl', 'F', 'Li', 'N', 'O', 'P', 'S', 'Si', '*']\n",
        "    \n",
        "    # Compute atom features\n",
        "    atom_type_enc = one_hot_encoding(str(atom.GetSymbol()), permitted_list_of_atoms)\n",
        "\n",
        "    n_hydrogens_enc = one_hot_encoding(int(atom.GetTotalNumHs()), [0, 1, 2, 3, 4, \"MoreThanFour\"])\n",
        "\n",
        "    implicit_valence_enc = one_hot_encoding(str(atom.GetImplicitValence()), [0, 1, 2, 3, 4, \"MoreThanFour\"])\n",
        "\n",
        "    degree_enc = one_hot_encoding(int(atom.GetDegree()), [0, 1, 2, 3, 4, \"MoreThanFour\"])\n",
        "\n",
        "    aromaticity_enc = [int(atom.GetIsAromatic())]\n",
        "\n",
        "    # Create atom feature vector\n",
        "    atom_feature_vector = atom_type_enc + n_hydrogens_enc + implicit_valence_enc + degree_enc + aromaticity_enc\n",
        "\n",
        "    return np.array(atom_feature_vector)"
      ],
      "metadata": {
        "id": "pw3eExGF8vK4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create Dataset"
      ],
      "metadata": {
        "id": "Dd3rtzqR-hTy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Subclass the Dataset class to standardize how graph datasets are represented in Spektral. Here, you can choose the endpoint to work with."
      ],
      "metadata": {
        "id": "qadRGli7V02I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MyDataset(Dataset):\n",
        "  def read(self):\n",
        "\n",
        "    global polymer_data # Access data globally\n",
        "    \n",
        "    # Set of 445 polymers with every endpoint (Tg, Tm, D) curated from Liu et al. (https://arxiv.org/pdf/2206.02886.pdf) \n",
        "    # polymer_data = pd.read_csv(\"https://raw.githubusercontent.com/carlosmatherson/PolymerGCN/main/SMILESwithAllProps.csv\")\n",
        "    # Label is \"Density\"\n",
        "    \n",
        "    # Full dataset from from Liu et al. (https://arxiv.org/pdf/2206.02886.pdf)\n",
        "    polymer_data = pd.read_csv(\"https://raw.githubusercontent.com/liugangcode/GREA/main/data/density_prop/raw/density_raw.csv\")\n",
        "    # Label is \"density\"\n",
        "    \n",
        "    dataset = makeGraphObjList(polymer_data[\"SMILES\"], polymer_data[\"density\"]) \n",
        "   \n",
        "    return dataset"
      ],
      "metadata": {
        "id": "zMcD28857JwV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define `makeGraphObjList` to create Spektral graph dataset from smiles and labels. The function takes a list of SMILES strings `x_smiles = [smiles_1, smiles_2, ....]` and a list of numerial labels for the SMILES strings `y = [y_1, y_2, ...]` as inputs. The output is a list of spektral.data.graph.Graph objects that can readily be used for machine learning with Spektral, `graph_list = [G_1, G_2, ...]`. \n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "hMFxmFiCy0jI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def makeGraphObjList(x_smiles, y):\n",
        "    \n",
        "    graph_list = []\n",
        "\n",
        "    for (smiles, y_val) in zip(x_smiles, y):\n",
        "        \n",
        "        # Convert SMILES to RDKit mol object\n",
        "        mol = Chem.MolFromSmiles(smiles)\n",
        "\n",
        "        # Avoid O=O\n",
        "        unrelated_smiles = \"O=O\"\n",
        "        unrelated_mol = Chem.MolFromSmiles(unrelated_smiles)\n",
        "        \n",
        "        # Get feature dimensions\n",
        "        n_nodes = mol.GetNumAtoms()\n",
        "        n_node_features = len(get_atom_features(unrelated_mol.GetAtomWithIdx(0)))\n",
        "\n",
        "        # Construct adjacency matrix A of shape(n_nodes, n_nodes)\n",
        "        A = np.array(GetAdjacencyMatrix(mol)) # full matrix\n",
        "\n",
        "        # Construct node feature matrix X of shape (n_nodes, n_node_features)\n",
        "        X = np.zeros((n_nodes, n_node_features))\n",
        "        for atom in mol.GetAtoms():\n",
        "            X[atom.GetIdx(), :] = get_atom_features(atom)\n",
        "                \n",
        "        # Construct endpoint tensor\n",
        "        Y = np.array([y_val])\n",
        "\n",
        "        # Construct Spektral graph object and append to dataset (list)\n",
        "        graph_list.append(Graph(x=X, a=A, e=None, y=Y))\n",
        "        \n",
        "    return graph_list "
      ],
      "metadata": {
        "id": "mUflMN2i-p7Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load data."
      ],
      "metadata": {
        "id": "lrbucS1G879u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data = MyDataset()  # Store data in data"
      ],
      "metadata": {
        "id": "49_LrRJM87ps"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Helper Functions"
      ],
      "metadata": {
        "id": "xuu4sTDFW4v3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define R-squared function not available through Keras."
      ],
      "metadata": {
        "id": "XQIkWsJ5EgKK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Coefficient of determination metric\n",
        "def coeff_determination(y_true, y_pred):\n",
        "  SS_res =  K.sum(K.square(y_true-y_pred))\n",
        "  SS_tot = K.sum(K.square(y_true - K.mean(y_true)))\n",
        "  return (1 - SS_res/(SS_tot + K.epsilon()))"
      ],
      "metadata": {
        "id": "p4ivqThKEftI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Graph Convolutional Network"
      ],
      "metadata": {
        "id": "LIQBZjmdZnay"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Configure hyperparameters."
      ],
      "metadata": {
        "id": "pjI97nIMZwRK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_split = 0.8                   # Train/Test split\n",
        "learning_rate = 1e-4                # Learning rate\n",
        "epochs = 4000                       # Number of training epochs, 4000\n",
        "batch_size = 15                     # Batch size \n",
        "\n",
        "num_folds = 5                       # Number of folds for Cross Validation\n",
        "optimizer = Adam(learning_rate)     # Optimizer\n",
        "l2 = regularizers.L2(1e-2)          # L2 kernel reg. with weight decay parameter\n",
        "validation_steps = 5                # Validation steps\n",
        "\n",
        "loss = 'mse'                        # Loss function\n",
        "metrics = [RootMeanSquaredError(), \n",
        "           coeff_determination]     # Metrics"
      ],
      "metadata": {
        "id": "z9_d_lqTDuQ_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define GCN architecture.\n"
      ],
      "metadata": {
        "id": "-gWGDAatgoxM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Model\n",
        "class myGCN(Model):\n",
        "\n",
        "    # Define GCN structure\n",
        "    def __init__(self):\n",
        "      super().__init__()\n",
        "      \n",
        "      # GCN+LR: 6 convolutional layers with 100 nodes each, batch input 0-padded\n",
        "      self.conv1 = GCNConv(100, activation=\"relu\", kernel_regularizer=l2)\n",
        "      self.batn1 = BatchNormalization()  \n",
        "      self.global_pool = GlobalSumPool()                        # graph embed\n",
        "      self.batn2 = BatchNormalization()\n",
        "      self.dense1 = Dense(data.n_labels, kernel_regularizer=l2) # GCN out\n",
        "      \n",
        "      # GCN+NN: 2 hidden layers with 300 nodes each\n",
        "      self.mlp1 = Dense(units=300, activation=\"relu\", kernel_regularizer=l2)\n",
        "      self.mlp2 = Dense(units=300, activation=\"relu\", kernel_regularizer=l2)\n",
        "      self.dense2 = Dense(data.n_labels, kernel_regularizer=l2) # output       \n",
        "\n",
        "\n",
        "    # Call GCN \n",
        "    def call(self, inputs):\n",
        "\n",
        "      # global variables for results\n",
        "      global embed, gcn_out, nn_out\n",
        "\n",
        "      # Calling GCN\n",
        "      x, a = inputs[0], inputs[1]\n",
        "      x = self.conv1([x, a])\n",
        "      x = self.batn1(x)\n",
        "      embed = x = self.global_pool(x)\n",
        "      x = self.batn2(x)\n",
        "      gcn_out = self.dense1(x)\n",
        "      x = self.mlp1(x)\n",
        "      x = self.mlp2(x)\n",
        "      nn_out = self.dense2(x)\n",
        "\n",
        "      # Return loss and metrics for GCN-LR and GCN-NN\n",
        "      return gcn_out, nn_out\n",
        "\n",
        "\n",
        "model = myGCN() # Compile model\n",
        "model.compile(optimizer=optimizer, metrics=metrics, loss=loss, run_eagerly=True)"
      ],
      "metadata": {
        "id": "8ZE87fbsvCM8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Training/Testing"
      ],
      "metadata": {
        "id": "p3OGweaqXGDC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Train and evaluate model with 5-fold cross validation"
      ],
      "metadata": {
        "id": "zKS9ESf0aksl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Implement 5-fold cross validation with shuffled dataset\n",
        "k = KFold(n_splits=num_folds, random_state=None, shuffle=True)\n",
        "\n",
        "loss = [] # array to store loss metrics\n",
        "i=0       # count to display loss information at end of loop\n",
        "\n",
        "# Iterate through folds, train data, evaluate model\n",
        "for train_index, test_index in k.split(data):\n",
        "  # Split data into training and testing sets\n",
        "  dataset_tr = data[train_index]  # Kth training set (fold)\n",
        "  dataset_te = data[test_index]   # Kth testing set (fold)\n",
        "  \n",
        "  # Load batch data from datasets with zero-padding and more shuffling\n",
        "  loader_tr = BatchLoader(dataset_tr, batch_size=batch_size, mask=True, shuffle=True)\n",
        "  loader_te = BatchLoader(dataset_te, batch_size=batch_size, mask=True, shuffle=True)\n",
        "  \n",
        "  # Fit model and store results\n",
        "  results = model.fit(loader_tr.load(), \n",
        "                      steps_per_epoch=loader_tr.steps_per_epoch, \n",
        "                      epochs=epochs,\n",
        "                      verbose=2\n",
        "                      )\n",
        "\n",
        "  print(\"\\nTesting model\") # Show test results for each fold\n",
        "  # Append loss metrics to array\n",
        "  loss.append(model.evaluate(loader_te.load(), steps=loader_te.steps_per_epoch))\n",
        "  \n",
        "  # Output test results\n",
        "  print(\"Done. \\n\\nTest loss: {}\".format(loss[i]))\n",
        "  print(\"\\n\")\n",
        "\n",
        "  i+=1 # increment counter for array\n",
        "  # End of loop\n"
      ],
      "metadata": {
        "id": "HW7N-re_bqU3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Results"
      ],
      "metadata": {
        "id": "2_rUQ8HQXN37"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create dataframe with results\n",
        "lossDF = pd.DataFrame(loss, columns = ['Tot. MSE','GCN MSE','NN MSE','GCN RMSE','GCN R^2','NN RMSE','NN R^2'])\n",
        "display(lossDF)\n",
        "lossDF.to_pickle('dGCN_performance.pkl')"
      ],
      "metadata": {
        "id": "-S3pikv45Fpt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Visualization (future work)"
      ],
      "metadata": {
        "id": "PwgqP7Gxv2HS"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Plot root mean square error per epoch."
      ],
      "metadata": {
        "id": "20EZLJ17ew46"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Plott RMSE per epoch\n",
        "fig, ax = plt.subplots()\n",
        "plt.plot(results.history['output_1_root_mean_squared_error'], label=\"GCN-LR\")\n",
        "plt.plot(results.history['output_2_root_mean_squared_error'], label=\"GCN-NN\")\n",
        "plt.title(\"Loss: GCN vs GCN-NN\")\n",
        "plt.xlabel(\"Epoch\")\n",
        "plt.ylabel(\"RMSE\")\n",
        "ax.legend()\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "ME-ODKthUvHV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Plot coefficient of determination."
      ],
      "metadata": {
        "id": "HKdv83FOe5jE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Plott R^2 "
      ],
      "metadata": {
        "id": "6A-jUmR-e9GK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Plot PCA "
      ],
      "metadata": {
        "id": "G_JuV1-ojoXr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a 2D PCA projection of the embeddings\n",
        "GCNrepresentation = embed.numpy()\n",
        "pca = PCA(n_components=2)\n",
        "pca.fit(GCNrepresentation)\n",
        "pca_gcn = pca.transform(GCNrepresentation)\n",
        "print(pca.explained_variance_ratio_)\n",
        "\n",
        "# Plot those points as a scatter plot \n",
        "fig, ax = plt.subplots(figsize=(8,8))\n",
        "num_categories = data.n_labels\n",
        "plt.title(\"PCA from Pooled Graph Embedding\")\n",
        "plt.xlabel(\"PC1\")\n",
        "plt.ylabel(\"PC2\")\n",
        "ax.legend()\n",
        "ax.scatter(pca_gcn[:,0],pca_gcn[:,1]) # c=np.array(cmap(lab)).reshape(1,4)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "cwdKu9wWet_G"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}