examples/tutorials/Learning_Unsupervised_Embeddings_for_Molecules.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hzpae9-r2aoK"
   },
   "source": [
    "#  Learning Unsupervised Embeddings for Molecules\n",
    "\n",
    "In this tutorial, we will use a `SeqToSeq` model to generate fingerprints for classifying molecules.  This is based on the following paper, although some of the implementation details are different: Xu et al., \"Seq2seq Fingerprint: An Unsupervised Deep Molecular Embedding for Drug Discovery\" (https://doi.org/10.1145/3107411.3107424).\n",
    "\n",
    "## Colab\n",
    "\n",
    "This tutorial and the rest in this sequence can be done in Google colab. If you'd like to open this notebook in colab, you can use the following link.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deepchem/deepchem/blob/master/examples/tutorials/Learning_Unsupervised_Embeddings_for_Molecules.ipynb)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 188
    },
    "colab_type": "code",
    "id": "2uo2i6arBiMS",
    "outputId": "d9d1d0ba-09c0-44ee-b315-84d87af40cf2"
   },
   "outputs": [],
   "source": [
    "!pip install --pre deepchem\n",
    "import deepchem\n",
    "deepchem.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6bm1iYbw2aoT"
   },
   "source": [
    "# Learning Embeddings with SeqToSeq\n",
    "\n",
    "Many types of models require their inputs to have a fixed shape.  Since molecules can vary widely in the numbers of atoms and bonds they contain, this makes it hard to apply those models to them.  We need a way of generating a fixed length \"fingerprint\" for each molecule.  Various ways of doing this have been designed, such as the Extended-Connectivity Fingerprints (ECFPs) we used in earlier tutorials.  But in this example, instead of designing a fingerprint by hand, we will let a `SeqToSeq` model learn its own method of creating fingerprints.\n",
    "\n",
    "A `SeqToSeq` model performs sequence to sequence translation.  For example, they are often used to translate text from one language to another.  It consists of two parts called the \"encoder\" and \"decoder\".  The encoder is a stack of recurrent layers.  The input sequence is fed into it, one token at a time, and it generates a fixed length vector called the \"embedding vector\".  The decoder is another stack of recurrent layers that performs the inverse operation: it takes the embedding vector as input, and generates the output sequence.  By training it on appropriately chosen input/output pairs, you can create a model that performs many sorts of transformations.\n",
    "\n",
    "In this case, we will use SMILES strings describing molecules as the input sequences.  We will train the model as an autoencoder, so it tries to make the output sequences identical to the input sequences.  For that to work, the encoder must create embedding vectors that contain all information from the original sequence.  That's exactly what we want in a fingerprint, so perhaps those embedding vectors will then be useful as a way to represent molecules in other models!\n",
    "\n",
    "Let's start by loading the data.  We will use the MUV dataset.  It includes 74,501 molecules in the training set, and 9313 molecules in the validation set, so it gives us plenty of SMILES strings to work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YnAnjl9d2aoU"
   },
   "outputs": [],
   "source": [
    "import deepchem as dc\n",
    "tasks, datasets, transformers = dc.molnet.load_muv(split='stratified')\n",
    "train_dataset, valid_dataset, test_dataset = datasets\n",
    "train_smiles = train_dataset.ids\n",
    "valid_smiles = valid_dataset.ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EslVHE2m2aoY"
   },
   "source": [
    "We need to define the \"alphabet\" for our `SeqToSeq` model, the list of all tokens that can appear in sequences.  (It's also possible for input and output sequences to have different alphabets, but since we're training it as an autoencoder, they're identical in this case.)  Make a list of every character that appears in any training sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nsE8e9xn2aoa"
   },
   "outputs": [],
   "source": [
    "tokens = set()\n",
    "for s in train_smiles:\n",
    "  tokens = tokens.union(set(c for c in s))\n",
    "tokens = sorted(list(tokens))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vgzyJ1-42aog"
   },
   "source": [
    "Create the model and define the optimization method to use.  In this case, learning works much better if we gradually decrease the learning rate.  We use an `ExponentialDecay` to multiply the learning rate by 0.9 after each epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NHKrymnM2aoh"
   },
   "outputs": [],
   "source": [
    "from deepchem.models.optimizers import Adam, ExponentialDecay\n",
    "max_length = max(len(s) for s in train_smiles)\n",
    "batch_size = 100\n",
    "batches_per_epoch = len(train_smiles)/batch_size\n",
    "model = dc.models.SeqToSeq(tokens,\n",
    "                           tokens,\n",
    "                           max_length,\n",
    "                           encoder_layers=2,\n",
    "                           decoder_layers=2,\n",
    "                           embedding_dimension=256,\n",
    "                           model_dir='fingerprint',\n",
    "                           batch_size=batch_size,\n",
    "                           learning_rate=ExponentialDecay(0.001, 0.9, batches_per_epoch))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hSr7FkSW2aok"
   },
   "source": [
    "Let's train it!  The input to `fit_sequences()` is a generator that produces input/output pairs.  On a good GPU, this should take a few hours or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NZ5l_g1E2aok"
   },
   "outputs": [],
   "source": [
    "def generate_sequences(epochs):\n",
    "  for i in range(epochs):\n",
    "    for s in train_smiles:\n",
    "      yield (s, s)\n",
    "\n",
    "model.fit_sequences(generate_sequences(40))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_lxf1lmX2aoo"
   },
   "source": [
    "Let's see how well it works as an autoencoder.  We'll run the first 500 molecules from the validation set through it, and see how many of them are exactly reproduced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NXDBtIvn2aop"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reproduced 161 of 500 validation SMILES strings\n"
     ]
    }
   ],
   "source": [
    "predicted = model.predict_from_sequences(valid_smiles[:500])\n",
    "count = 0\n",
    "for s,p in zip(valid_smiles[:500], predicted):\n",
    "  if ''.join(p) == s:\n",
    "    count += 1\n",
    "print('reproduced', count, 'of 500 validation SMILES strings')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Rt9GLy502aou"
   },
   "source": [
    "Now we'll trying using the encoder as a way to generate molecular fingerprints.  We compute the embedding vectors for all molecules in the training and validation datasets, and create new datasets that have those as their feature vectors.  The amount of data is small enough that we can just store everything in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kdUfsbtZ2aov"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "train_embeddings = model.predict_embeddings(train_smiles)\n",
    "train_embeddings_dataset = dc.data.NumpyDataset(train_embeddings,\n",
    "                                                train_dataset.y,\n",
    "                                                train_dataset.w.astype(np.float32),\n",
    "                                                train_dataset.ids)\n",
    "\n",
    "valid_embeddings = model.predict_embeddings(valid_smiles)\n",
    "valid_embeddings_dataset = dc.data.NumpyDataset(valid_embeddings,\n",
    "                                                valid_dataset.y,\n",
    "                                                valid_dataset.w.astype(np.float32),\n",
    "                                                valid_dataset.ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lVvfGr562aoz"
   },
   "source": [
    "For classification, we'll use a simple fully connected network with one hidden layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tFmnnVNm2aoz"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0014195525646209716"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier = dc.models.MultitaskClassifier(n_tasks=len(tasks),\n",
    "                                                      n_features=256,\n",
    "                                                      layer_sizes=[512])\n",
    "classifier.fit(train_embeddings_dataset, nb_epoch=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "khdB2v7R2ao2"
   },
   "source": [
    "Find out how well it worked.  Compute the ROC AUC for the training and validation datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZlilhPvm2ao2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training set ROC AUC: {'mean-roc_auc_score': 0.9598792603154332}\n",
      "Validation set ROC AUC: {'mean-roc_auc_score': 0.7251350862464794}\n"
     ]
    }
   ],
   "source": [
    "metric = dc.metrics.Metric(dc.metrics.roc_auc_score, np.mean, mode=\"classification\")\n",
    "train_score = classifier.evaluate(train_embeddings_dataset, [metric], transformers)\n",
    "valid_score = classifier.evaluate(valid_embeddings_dataset, [metric], transformers)\n",
    "print('Training set ROC AUC:', train_score)\n",
    "print('Validation set ROC AUC:', valid_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ixqbRXnW2ao6"
   },
   "source": [
    "# Congratulations! Time to join the Community!\n",
    "\n",
    "Congratulations on completing this tutorial notebook! If you enjoyed working through the tutorial, and want to continue working with DeepChem, we encourage you to finish the rest of the tutorials in this series. You can also help the DeepChem community in the following ways:\n",
    "\n",
    "## Star DeepChem on [GitHub](https://github.com/deepchem/deepchem)\n",
    "This helps build awareness of the DeepChem project and the tools for open source drug discovery that we're trying to build.\n",
    "\n",
    "## Join the DeepChem Gitter\n",
    "The DeepChem [Gitter](https://gitter.im/deepchem/Lobby) hosts a number of scientists, developers, and enthusiasts interested in deep learning for the life sciences. Join the conversation!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "11_Learning_Unsupervised_Embeddings_for_Molecules.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}