Copy_of_dna_sequencing_error_correction.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of dna_sequencing_error_correction.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Ebenco36/user-management/blob/master/Copy_of_dna_sequencing_error_correction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gg0TesVpTrNO"
      },
      "source": [
        "# Using Nucleus and TensorFlow for DNA Sequencing Error Correction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k8dIGDxhyWX4"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/nucleus/blob/master/nucleus/examples/dna_sequencing_error_correction.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/nucleus/blob/master/nucleus/examples/dna_sequencing_error_correction.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "VyiSHJ_TN4lp"
      },
      "source": [
        "# @markdown Copyright 2019 Google LLC. \\\n",
        "# @markdown SPDX-License-Identifier: Apache-2.0\n",
        "# @markdown (license hidden in Colab)\n",
        "\n",
        "# Copyright 2019 Google LLC\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_FozdnuROr1i"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "In this tutorial, we formulate DNA sequencing error correction as a multiclass classification problem and propose two deep learning solutions. Our first approach corrects errors in a single read, whereas the second approach, shown in Figure 1, builds a consensus from several reads to predict the correct DNA sequence. We implement the second approach using the [Nucleus](https://github.com/google/nucleus) and [TensorFlow](https://www.tensorflow.org/) libraries. Our goal is to show how Nucleus can be used alongside TensorFlow for solving machine learning problems in genomics.\n",
        "\n",
        "<center>\n",
        "<img src='https://github.com/google/nucleus/raw/master/nucleus/examples/images/consensus-approach-overview.jpg' width='900'/>\n",
        "</center>\n",
        "\n",
        "\n",
        "### Problem Overview\n",
        "\n",
        "While DNA sequencing continues to become faster and cheaper, it is still an error-prone process. Error rates for raw data from next-generation sequencing (NGS) technologies developed by companies such as [Illumina](https://www.illumina.com/) are around 1%. Error rates for increasingly popular third-generation technologies like those developed by [Pacific BioSciences](https://www.pacb.com/) (PacBio) are around 15%. Sequencing errors can be divided into substitutions, insertions, and deletions, the last two of which are commonly referred to as indels.  All of these errors can be detrimental to downstream analysis steps such as variant calling and genome assembly.\n",
        "\n",
        "A simple approach for obtaining higher quality datasets is to discard data that likely contains errors, either by throwing away entire reads or trimming regions of low quality. This approach is not ideal as it leads to a smaller final dataset. In addition, certain sequence contexts have naturally higher error rates, leading to biases in sampling. Thus, there exists a large body of research that is focused on developing more sophisticated methods for error correction. Most methods that have been developed can be categorized into one of two groups:\n",
        "\n",
        "1. Methods that operate on a single read and aim to determine the correct read sequence\n",
        "1. Consensus-based methods that operate on several reads and aim to determine the correct underlying DNA sequence\n",
        "\n",
        "\n",
        "### Deep Learning Overview\n",
        "\n",
        "Both of the methods that we formulate in this post use deep neural networks, which learn functions that map inputs to outputs. A neural network consists of several layers of linear and nonlinear operations applied sequentially to the input. Neural networks have been successfully applied to various problems including [image classification](https://ai.googleblog.com/2016/03/train-your-own-image-classifier-with.html) and [natural language translation](https://ai.googleblog.com/2016/09/a-neural-network-for-machine.html). More recently, they have also been used for problems in genomics, such as [protein structure prediction](https://deepmind.com/blog/alphafold/) and [variant calling](https://ai.googleblog.com/2017/12/deepvariant-highly-accurate-genomes.html).\n",
        "\n",
        "### Nucleus\n",
        "\n",
        "Our implementation relies on [Nucleus](https://github.com/google/nucleus), a library developed for processing genomics data by the Genomics team in Google Brain. Nucleus makes it easy to read, write, and analyze data in common genomics file formats like BAM, FASTA, and VCF using specialized reader and writer objects. Nucleus allows us to:\n",
        "\n",
        "* Query a VCF file for all variants in a given genomic region\n",
        "* Query a BAM file for all reads mapping to a given genomic range\n",
        "* Query a FASTA file for the reference sequence starting at a given position\n",
        "\n",
        "We also use Nucleus to write data out to [TFRecords](https://www.tensorflow.org/tutorials/load_data/tf_records), a binary file format that consists of protocol buffer messages and can be easily read by TensorFlow. After reading in the TFRecords, we use the [TensorFlow Keras API](https://www.tensorflow.org/api_docs/python/tf/keras) to train and evaluate a convolutional neural network.\n",
        "\n",
        "### Data\n",
        "\n",
        "Below is a list of the files we use in the implementation. All of the data is publicly available, and the Appendix contains download links and instructions.\n",
        "\n",
        "File | Description\n",
        "--- | ---\n",
        "`NA12878_sliced.bam` | Illumina HiSeq reads from chromosome 20 (positions 10,000,000-10,100,000), downsampled to 30x coverage.\n",
        "`NA12878_sliced.bam.bai` | Index for `NA12878_sliced.bam`.\n",
        "`NA12878_calls.vcf.gz`| Truth set of variants for NA12878 from Genome in a Bottle.\n",
        "`NA12878_calls.vcf.gz.tbi` | Index for `NA12878_calls.vcf.gz`.\n",
        "`hs37d5.fa.gz` | Reference genome for hs37d5.\n",
        "`hs37d5.fa.gz.fai` and  `hs37d5.fa.gz.gzi` | Index files for `hs37d5.fa.gz`.\n",
        "\n",
        "### Questions or Comments?\n",
        "\n",
        "If you have any questions or comments regarding this tutorial, do not hesitate to reach out! You can [file an issue](https://github.com/google/nucleus/issues/new) on the Nucleus GitHub page.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vqusDUt3YEnB"
      },
      "source": [
        "## Setup\n",
        "\n",
        "If you are new to Colab or Jupyter notebooks, we recommend that you first go through this [tutorial](https://colab.research.google.com/notebooks/basic_features_overview.ipynb).\n",
        "\n",
        "### Obtain Data, Install Nucleus, and Import Packages\n",
        "\n",
        "Run the below cells to obtain the data, install Nucleus, and import Python packages.\n",
        "\n",
        "Note, the code for some cells is hidden for clarity. These cells are marked with the following text: `(code hidden in Colab)`. If you wish to view the code for a hidden cell, double click the cell. To hide the code, double click the markdown output on the right side."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "WrFYMMOlNS9s"
      },
      "source": [
        "#@markdown Run this cell to obtain the data. (code hidden in Colab)\n",
        "\n",
        "%%capture\n",
        "!gsutil cp gs://deepvariant/case-study-testdata/NA12878_sliced.bam NA12878_sliced.bam\n",
        "!gsutil cp gs://deepvariant/case-study-testdata/NA12878_sliced.bam.bai NA12878_sliced.bam.bai\n",
        "\n",
        "!wget ftp://ftp-trace.ncbi.nlm.nih.gov/giab/ftp/release/NA12878_HG001/latest/GRCh37/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer.vcf.gz -O NA12878_calls.vcf.gz\n",
        "!wget ftp://ftp-trace.ncbi.nlm.nih.gov/giab/ftp/release/NA12878_HG001/latest/GRCh37/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer.vcf.gz.tbi -O NA12878_calls.vcf.gz.tbi\n",
        "\n",
        "!gsutil cp gs://deepvariant/case-study-testdata/hs37d5.fa.gz hs37d5.fa.gz\n",
        "!gsutil cp gs://deepvariant/case-study-testdata/hs37d5.fa.gz.fai hs37d5.fa.gz.fai\n",
        "!gsutil cp gs://deepvariant/case-study-testdata/hs37d5.fa.gz.gzi hs37d5.fa.gz.gzi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "_OyDmvs5zY36"
      },
      "source": [
        "#@markdown Run this cell to install Nucleus and TensorFlow 2.0. (code hidden in Colab)\n",
        "\n",
        "!pip install -q google-nucleus==0.2.2\n",
        "!pip install -q tensorflow==2.0.0-alpha0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KkAxQO5e8q5L"
      },
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import os\n",
        "import random\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras import layers\n",
        "\n",
        "from nucleus.io import fasta\n",
        "from nucleus.io import sam\n",
        "from nucleus.io import vcf\n",
        "from nucleus.io.genomics_writer import TFRecordWriter\n",
        "from nucleus.protos import reads_pb2\n",
        "from nucleus.util import cigar\n",
        "from nucleus.util import ranges\n",
        "from nucleus.util import utils"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oPfWS6iLavkx"
      },
      "source": [
        "# Define constants and utility functions.\n",
        "\n",
        "# We will only allow simple alignments, specified by the below cigar string\n",
        "# operators. If you are not familiar with cigar strings, you can read more\n",
        "# at section 1.4 of this link: https://samtools.github.io/hts-specs/SAMv1.pdf\n",
        "_ALLOWED_CIGAR_OPS = frozenset([cigar.CHAR_TO_CIGAR_OPS[op] for op in 'MX='])\n",
        "\n",
        "# We will only allow certain bases.\n",
        "_ALLOWED_BASES = 'ACGT'\n",
        "\n",
        "_TRAIN = 'train.tfrecord'\n",
        "_EVAL = 'eval.tfrecord'\n",
        "_TEST = 'test.tfrecord'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-bfPoKedLQkz"
      },
      "source": [
        "## Network Architecture\n",
        "\n",
        "Convolutional neural networks are commonly used for computer vision tasks, but also [work well for genomics](https://doi.org/10.1093/molbev/msy224). Each convolutional layer repeatedly applies learned filters to the input. Convolutional filters appearing early in the network learn to recognize low-level features in the input, like edges and color gradients in images, whereas later filters learn to recognize more complex compositions of the low-level features. For DNA sequence inputs, low-level convolutional filters act as motif detectors, similar to the position weight matrices of [sequence logos](https://en.wikipedia.org/wiki/Position_weight_matrix#/media/File:LexA_gram_positive_bacteria_sequence_logo.png).\n",
        "\n",
        "For our implementation, we use a standard convolutional architecture consisting of two convolutional layers, followed by three fully connected layers. We use nonlinear ReLU layers to increase the expressive capacity of our model. Maxpooling after convolutional layers shrinks the input volume, and dropout after fully connected layers acts as a regularizer. The final softmax layer normalizes the logits to produce a valid probability distribution. The details of each layer can be found in the code below.\n",
        "\n",
        "<center>\n",
        "<img src='https://user-images.githubusercontent.com/8205702/56754313-8c8b8600-6741-11e9-9d8b-c9f547f158ba.jpg' width='900'/>\n",
        "</center>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pT-7ed7l-Mqa"
      },
      "source": [
        "def build_model(hparams):\n",
        "  \"\"\"Convolutional neural network architecture.\"\"\"\n",
        "\n",
        "  l2_reg = tf.keras.regularizers.l2\n",
        "\n",
        "  return tf.keras.models.Sequential([\n",
        "\n",
        "      # Two convolution + maxpooling blocks\n",
        "      layers.Conv1D(\n",
        "          filters=16,\n",
        "          kernel_size=5,\n",
        "          activation=tf.nn.relu,\n",
        "          kernel_regularizer=l2_reg(hparams.l2)),\n",
        "      layers.MaxPool1D(pool_size=3, strides=1),\n",
        "      layers.Conv1D(\n",
        "          filters=16,\n",
        "          kernel_size=3,\n",
        "          activation=tf.nn.relu,\n",
        "          kernel_regularizer=l2_reg(hparams.l2)),\n",
        "      layers.MaxPool1D(pool_size=3, strides=1),\n",
        "\n",
        "      # Flatten the input volume\n",
        "      layers.Flatten(),\n",
        "\n",
        "      # Two fully connected layers, each followed by a dropout layer\n",
        "      layers.Dense(\n",
        "          units=16,\n",
        "          activation=tf.nn.relu,\n",
        "          kernel_regularizer=l2_reg(hparams.l2)),\n",
        "      layers.Dropout(rate=0.3),\n",
        "      layers.Dense(\n",
        "          units=16,\n",
        "          activation=tf.nn.relu,\n",
        "          kernel_regularizer=l2_reg(hparams.l2)),\n",
        "      layers.Dropout(rate=0.3),\n",
        "\n",
        "      # Output layer with softmax activation\n",
        "      layers.Dense(units=len(_ALLOWED_BASES), activation='softmax')\n",
        "  ])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tI6JPQBw7cOm"
      },
      "source": [
        "## Approach 1: Error Correction of Single Reads\n",
        "\n",
        "In order to correct errors in sequenced reads, we can use deep learning to train a neural network that can solve a more general task: fill in missing bases in DNA sequences. The goal of this approach is to develop a model that understands the grammar of DNA sequences. The grammar of real sequences alone likely does not contain enough information to develop a solution that can be used in production. Nonetheless, this serves as a straightforward example application.\n",
        "\n",
        "<center>\n",
        "<img src='https://github.com/google/nucleus/raw/master/nucleus/examples/images/single-read-input-and-output.jpg' width='900'/>\n",
        "</center>\n",
        "\n",
        "For instructional purposes, we simplify the problem in the following ways:\n",
        "* Consider only regions with substitution errors and ignore indel errors\n",
        "* Consider only regions with no known variants\n",
        "\n",
        "We can train the neural network on regions of the reference genome. The input to this network is a DNA sequence of fixed length, centered around the base we wish to predict. The output of the network is a distribution over the possible bases, and the final prediction is the base with highest probability. The label set is generated using the bases observed in the reference genome. Since we only use reads mapping to regions with no known truth variants, we can unambiguously denote the base present in the reference genome as the label.\n",
        "\n",
        "We generate input sequences by splitting the reference genome into non-overlapping sections of a fixed length. At training, evaluation, and test time, we simulate missing bases by zeroing out a base in the reference sequence, as shown in Figure 3 (position 5). In addition to simulating missing data using the reference genome, we can also apply such a model to data from sequenced reads, specifically bases with quality scores below a threshold value.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3_XD6iRP7nO8"
      },
      "source": [
        "## Approach 2: Consensus-Based Error Correction\n",
        "\n",
        "The ultimate goal of error correction is to determine the underlying DNA sequence, as opposed to correcting an individual read. In this section, we use the consensus of multiple reads by aggregating a sequence pileup to directly determine the DNA sequence without the intermediate step of correcting individual reads. An example of a pileup is shown below in Figure 4. Note, the figure only shows the portions of the reads that fall in the window.\n",
        "\n",
        "<center>\n",
        "<img src='https://github.com/google/nucleus/raw/master/nucleus/examples/images/read-pileup.jpg' width='700'/>\n",
        "</center>\n",
        "\n",
        "For instructional purposes, we again simplify the problem in the following ways:\n",
        "* Consider only regions with substitution errors and ignore indel errors\n",
        "* Consider only regions with no known variants\n",
        "\n",
        "<center>\n",
        "<img src='https://github.com/google/nucleus/raw/master/nucleus/examples/images/raw-counts.jpg' width='700'/>\n",
        "</center>\n",
        "\n",
        "Unlike the first approach, we do not train this model on the reference genome. Instead, our training data comes from mapped Illumina HiSeq reads. The input to this network is a matrix of normalized base counts observed in mapped reads, centered around the position at which we wish to predict the correct base. A similar featurization is used by the authors of [Clairvoyante](https://doi.org/10.1101/310458), a neural network for variant calling, and in an [example method by Jason Chin](https://towardsdatascience.com/simple-convolution-neural-network-for-genomic-variant-calling-with-tensorflow-c085dbc2026f). The output of the network is a distribution over the possible bases, and the final prediction is the base with highest probability. Similar to the first approach, the label set is generated using the bases observed in the reference genome. We use a mix of examples that contain errors (at least one read in the pileup does not match the reference at the center position) and examples that do not contain errors (all reads in the pileup match the reference at the center position).\n",
        "\n",
        "<center>\n",
        "<img src='https://github.com/google/nucleus/raw/master/nucleus/examples/images/consensus-input-and-output.jpg' width='900'/>\n",
        "</center>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p1W1RY1Z9tny"
      },
      "source": [
        "### Data Processing\n",
        "\n",
        "The data for this problem comes from 148bp mapped Illumina HiSeq reads. In processing the data, we compare each read to the reference sequence, and any positions that differ from the reference are denoted as errors.\n",
        "\n",
        "* For reads containing one or more errors, we randomly choose an error and create one example centered at the corresponding position.\n",
        "* For reads containing no errors, we create one example centered at the middle position.\n",
        "\n",
        "Once we have determined the genomic window for the example, we use Nucleus to query for all reads mapping to the window. We then build a normalized counts matrix, as shown above. For some of the reads, only a subset of all bases will fall inside the window, and we ignore the bases that fall outside the window."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "COc_1OeL-PYg"
      },
      "source": [
        "### Implement Neural Network Pipeline\n",
        "\n",
        "We divide up our pipeline into the following steps, for each of which we implement several functions.\n",
        "\n",
        "1. [Generate TFRecords datasets](#scrollTo=ZlCQkTr-u-9Y)\n",
        "1. [Read data from TFRecords datasets](#scrollTo=sQ1-zPVBvIr0)\n",
        "1. [Train the model](#scrollTo=tgiVn4dZvdLs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZlCQkTr-u-9Y"
      },
      "source": [
        "**Step 1: Generate TFRecords Datasets**\n",
        "\n",
        "We generate TFRecords datasets for training, evaluation, and testing. All examples that do not meet the criteria expressed in `is_usable_example` are discarded."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OEvFDwGFG-MJ"
      },
      "source": [
        "def generate_tfrecord_datasets(hparams):\n",
        "  \"\"\"Writes out TFRecords files for training, evaluation, and test datasets.\"\"\"\n",
        "  if not os.path.exists(hparams.out_dir):\n",
        "    os.makedirs(hparams.out_dir)\n",
        "\n",
        "  # Fraction of examples in each dataset.\n",
        "  train_eval_test_split = [0.7, 0.2, 0.1]\n",
        "  num_train_examples = 0\n",
        "  num_eval_examples = 0\n",
        "  num_test_examples = 0\n",
        "\n",
        "  # Generate training, test, and evaluation examples.\n",
        "  with TFRecordWriter(os.path.join(hparams.out_dir, _TRAIN)) as train_out, \\\n",
        "       TFRecordWriter(os.path.join(hparams.out_dir, _EVAL)) as eval_out, \\\n",
        "       TFRecordWriter(os.path.join(hparams.out_dir, _TEST)) as test_out:\n",
        "    all_examples = make_ngs_examples(hparams)\n",
        "    for example in all_examples:\n",
        "      r = random.random()\n",
        "      if r < train_eval_test_split[0]:\n",
        "        train_out.write(proto=example)\n",
        "        num_train_examples += 1\n",
        "      elif r < train_eval_test_split[0] + train_eval_test_split[1]:\n",
        "        eval_out.write(proto=example)\n",
        "        num_eval_examples += 1\n",
        "      else:\n",
        "        test_out.write(proto=example)\n",
        "        num_test_examples += 1\n",
        "  print('# of training examples: %d' % num_train_examples)\n",
        "  print('# of evaluation examples: %d' % num_eval_examples)\n",
        "  print('# of test examples: %d' % num_test_examples)\n",
        "\n",
        "\n",
        "def make_ngs_examples(hparams):\n",
        "  \"\"\"Generator function that yields training, evaluation and test examples.\"\"\"\n",
        "  ref_reader = fasta.IndexedFastaReader(input_path=hparams.ref_path)\n",
        "  vcf_reader = vcf.VcfReader(input_path=hparams.vcf_path)\n",
        "  read_requirements = reads_pb2.ReadRequirements()\n",
        "  sam_reader = sam.SamReader(\n",
        "      input_path=hparams.bam_path, read_requirements=read_requirements)\n",
        "\n",
        "  # Use a separate SAM reader to query for reads falling in the pileup range.\n",
        "  sam_query_reader = sam.SamReader(\n",
        "      input_path=hparams.bam_path, read_requirements=read_requirements)\n",
        "  used_pileup_ranges = set()\n",
        "  with ref_reader, vcf_reader, sam_reader, sam_query_reader:\n",
        "    for read in sam_reader:\n",
        "\n",
        "      # Check that read has cigar string present and allowed alignment.\n",
        "      if not read.alignment.cigar:\n",
        "        print('Skipping read, no cigar alignment found')\n",
        "        continue\n",
        "      if not has_allowed_alignment(read):\n",
        "        continue\n",
        "\n",
        "      # Obtain window that will be used to construct an example.\n",
        "      read_range = utils.read_range(read)\n",
        "      ref = ref_reader.query(region=read_range)\n",
        "      pileup_range = get_pileup_range(hparams, read, read_range, ref)\n",
        "\n",
        "      # Do not construct multiple examples with the same pileup range.\n",
        "      pileup_range_serialized = pileup_range.SerializeToString()\n",
        "      if pileup_range_serialized in used_pileup_ranges:\n",
        "        continue\n",
        "      used_pileup_ranges.add(pileup_range_serialized)\n",
        "\n",
        "      # Get reference sequence, reads, and truth variants for the pileup range.\n",
        "      pileup_reads = list(sam_query_reader.query(region=pileup_range))\n",
        "      pileup_ref = ref_reader.query(region=pileup_range)\n",
        "      pileup_variants = list(vcf_reader.query(region=pileup_range))\n",
        "      if is_usable_example(pileup_reads, pileup_variants, pileup_ref):\n",
        "        yield make_example(hparams, pileup_reads, pileup_ref, pileup_range)\n",
        "\n",
        "\n",
        "def get_pileup_range(hparams, read, read_range, ref):\n",
        "  \"\"\"Returns a range that will be used to construct one example.\"\"\"\n",
        "\n",
        "  # Find error positions where read and reference differ.\n",
        "  ngs_read_length = read_range.end - read_range.start\n",
        "  error_indices = [\n",
        "      i for i in range(ngs_read_length) if ref[i] != read.aligned_sequence[i]\n",
        "  ]\n",
        "\n",
        "  # If read and reference sequence are the same, create an example centered\n",
        "  # at middle base of read.\n",
        "  if not error_indices:\n",
        "    error_idx = ngs_read_length // 2\n",
        "\n",
        "  # If read and reference differ at one or more positions, create example\n",
        "  # centered at a random error position.\n",
        "  else:\n",
        "    error_idx = random.choice(error_indices)\n",
        "\n",
        "  error_pos = read_range.start + error_idx\n",
        "  flank_size = hparams.window_size // 2\n",
        "  return ranges.make_range(\n",
        "      chrom=read_range.reference_name,\n",
        "      start=error_pos - flank_size,\n",
        "      end=error_pos + flank_size + 1)\n",
        "\n",
        "\n",
        "def has_allowed_alignment(read):\n",
        "  \"\"\"Determines whether a read's CIGAR string has the allowed alignments.\"\"\"\n",
        "  return all([c.operation in _ALLOWED_CIGAR_OPS for c in read.alignment.cigar])\n",
        "\n",
        "\n",
        "def is_usable_example(reads, variants, ref_bases):\n",
        "  \"\"\"Determines whether a particular reference region and read can be used.\"\"\"\n",
        "  # Discard examples with variants or no mapped reads.\n",
        "  if variants or not reads:\n",
        "    return False\n",
        "\n",
        "  # Use only examples where all reads have simple alignment and allowed bases.\n",
        "  for read in reads:\n",
        "    if not has_allowed_alignment(read):\n",
        "      return False\n",
        "    if any(base not in _ALLOWED_BASES for base in read.aligned_sequence):\n",
        "      return False\n",
        "\n",
        "  # Reference should only contain allowed bases.\n",
        "  if any(base not in _ALLOWED_BASES for base in ref_bases):\n",
        "    return False\n",
        "  return True\n",
        "\n",
        "\n",
        "def make_example(hparams, pileup_reads, pileup_ref, pileup_range):\n",
        "  \"\"\"Takes in an input sequence and outputs tf.train.Example ProtocolMessages.\n",
        "\n",
        "  Each example contains the following features: A counts, C counts, G counts,\n",
        "  T counts, reference sequence, correct base label.\n",
        "  \"\"\"\n",
        "  assert len(pileup_ref) == hparams.window_size\n",
        "  example = tf.train.Example()\n",
        "  base_counts = np.zeros(shape=[hparams.window_size, len(_ALLOWED_BASES)])\n",
        "\n",
        "  for read in pileup_reads:\n",
        "    read_position = read.alignment.position.position\n",
        "    read_ints = [_ALLOWED_BASES.index(b) for b in read.aligned_sequence]\n",
        "    one_hot_read = np.zeros((len(read_ints), len(_ALLOWED_BASES)))\n",
        "    one_hot_read[np.arange(len(one_hot_read)), read_ints] = 1\n",
        "\n",
        "    window_start = read_position - pileup_range.start\n",
        "    window_end = window_start + len(read_ints)\n",
        "\n",
        "    # If read falls outside of window, adjust start/end indices for window.\n",
        "    window_start = max(0, window_start)\n",
        "    window_end = min(window_end, hparams.window_size)\n",
        "\n",
        "    # We consider four possible scenarios for each read and adjust start/end\n",
        "    # indices to only include portions of read that overlap the window.\n",
        "    # 1) Read extends past 5' end of window\n",
        "    # 2) Read extends past 3' end of window\n",
        "    # 3) Read extends past 5' and 3' ends of window\n",
        "    # 4) Read falls entirely within window\n",
        "    if window_start == 0 and window_end != hparams.window_size:\n",
        "      read_start = pileup_range.start - read_position\n",
        "      read_end = None\n",
        "    if window_end == hparams.window_size and window_start != 0:\n",
        "      read_start = None\n",
        "      read_end = -1 * ((read_position + len(read_ints)) - pileup_range.end)\n",
        "    if window_start == 0 and window_end == hparams.window_size:\n",
        "      read_start = pileup_range.start - read_position\n",
        "      read_end = read_start + hparams.window_size\n",
        "    if window_start != 0 and window_end != hparams.window_size:\n",
        "      read_start = None\n",
        "      read_end = None\n",
        "    base_counts[window_start:window_end] += one_hot_read[read_start:read_end]\n",
        "\n",
        "  # Use fractions at each position instead of raw base counts.\n",
        "  base_counts /= np.expand_dims(np.sum(base_counts, axis=-1), -1)\n",
        "\n",
        "  # Save counts/fractions for each base separately.\n",
        "  features = example.features\n",
        "  for i in range(len(_ALLOWED_BASES)):\n",
        "    key = '%s_counts' % _ALLOWED_BASES[i]\n",
        "    features.feature[key].float_list.value.extend(list(base_counts[:, i]))\n",
        "\n",
        "  features.feature['ref_sequence'].int64_list.value.extend(\n",
        "      [_ALLOWED_BASES.index(base) for base in pileup_ref])\n",
        "  flank_size = hparams.window_size // 2\n",
        "  true_base = pileup_ref[flank_size]\n",
        "  features.feature['label'].int64_list.value.append(\n",
        "      _ALLOWED_BASES.index(true_base))\n",
        "\n",
        "  return example"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sQ1-zPVBvIr0"
      },
      "source": [
        "**Step 2: Read TFRecords Datasets**\n",
        "\n",
        "We define an input function to read in the TFRecords dataset. Since TFRecords is an unstructured binary format, it is necessary to define the structure of the data in order to read it back in. Specifically, we must define the type and size of each field in proto_features."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "code",
        "id": "zVc6yRIjQoRb"
      },
      "source": [
        "def get_dataset(hparams, filename, num_epochs):\n",
        "  \"\"\"Reads in and processes the TFRecords dataset.\n",
        "\n",
        "  Builds a pipeline that returns pairs of features, label.\n",
        "  \"\"\"\n",
        "\n",
        "  # Define field names, types, and sizes for TFRecords.\n",
        "  proto_features = {\n",
        "      'A_counts':\n",
        "          tf.io.FixedLenFeature(shape=[hparams.window_size], dtype=tf.float32),\n",
        "      'C_counts':\n",
        "          tf.io.FixedLenFeature(shape=[hparams.window_size], dtype=tf.float32),\n",
        "      'G_counts':\n",
        "          tf.io.FixedLenFeature(shape=[hparams.window_size], dtype=tf.float32),\n",
        "      'T_counts':\n",
        "          tf.io.FixedLenFeature(shape=[hparams.window_size], dtype=tf.float32),\n",
        "      'ref_sequence':\n",
        "          tf.io.FixedLenFeature(shape=[hparams.window_size], dtype=tf.int64),\n",
        "      'label':\n",
        "          tf.io.FixedLenFeature(shape=[1], dtype=tf.int64),\n",
        "  }\n",
        "\n",
        "  def _process_input(proto_string):\n",
        "    \"\"\"Helper function for input function that parses a serialized example.\"\"\"\n",
        "\n",
        "    parsed_features = tf.io.parse_single_example(\n",
        "        serialized=proto_string, features=proto_features)\n",
        "\n",
        "    # Stack counts/fractions for all bases to create input of dimensions\n",
        "    # `hparams.window_size` x len(_ALLOWED_BASES).\n",
        "    feature_columns = []\n",
        "    for base in _ALLOWED_BASES:\n",
        "      feature_columns.append(parsed_features['%s_counts' % base])\n",
        "    features = tf.stack(feature_columns, axis=-1)\n",
        "    label = parsed_features['label']\n",
        "    return features, label\n",
        "\n",
        "  ds = tf.data.TFRecordDataset(filenames=filename)\n",
        "  ds = ds.map(map_func=_process_input)\n",
        "  ds = ds.shuffle(buffer_size=10000, reshuffle_each_iteration=True)\n",
        "  ds = ds.batch(batch_size=hparams.batch_size).repeat(count=num_epochs)\n",
        "  return ds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tgiVn4dZvdLs"
      },
      "source": [
        "**Step 3: Train the Model**\n",
        "\n",
        "We use `tf.keras` to train and evaluate our model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Skgb4NFvavF"
      },
      "source": [
        "def run(hparams, use_existing_data=False, seed=1):\n",
        "  \"\"\"Creates a model, runs training and evaluation.\"\"\"\n",
        "\n",
        "  # Set seed for reproducibility.\n",
        "  random.seed(seed)\n",
        "  tf.random.set_seed(seed)\n",
        "\n",
        "  if not use_existing_data:\n",
        "    print('Generating data...')\n",
        "    generate_tfrecord_datasets(hparams)\n",
        "\n",
        "  train_dataset = get_dataset(\n",
        "      hparams, filename=os.path.join(hparams.out_dir, _TRAIN), num_epochs=1)\n",
        "  eval_dataset = get_dataset(\n",
        "      hparams, filename=os.path.join(hparams.out_dir, _EVAL), num_epochs=1)\n",
        "  test_dataset = get_dataset(\n",
        "      hparams, filename=os.path.join(hparams.out_dir, _TEST), num_epochs=1)\n",
        "\n",
        "  optimizer = tf.keras.optimizers.Adam(lr=hparams.learning_rate)\n",
        "  tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
        "      hparams.log_dir, histogram_freq=1)\n",
        "  model = build_model(hparams)\n",
        "  model.compile(\n",
        "      optimizer=optimizer,\n",
        "      loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
        "      metrics=['accuracy'])\n",
        "\n",
        "  print('Training the model. This should take ~6 minutes...')\n",
        "  model.fit_generator(\n",
        "      train_dataset,\n",
        "      epochs=hparams.total_epochs,\n",
        "      validation_data=eval_dataset,\n",
        "      callbacks=[tensorboard_callback],\n",
        "      verbose=0)\n",
        "  print('Training complete. Obtaining final metrics...')\n",
        "  eval_metrics = model.evaluate(eval_dataset, verbose=0)\n",
        "  test_metrics = model.evaluate(test_dataset, verbose=0)\n",
        "  print('Final eval metrics - loss: %f - accuracy: %f' %\n",
        "        (eval_metrics[0], eval_metrics[1]))\n",
        "  print('Final test metrics - loss: %f - accuracy: %f' %\n",
        "        (test_metrics[0], test_metrics[1]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vqk2BEXxemv9"
      },
      "source": [
        "### Train and Evaluate Neural Network\n",
        "\n",
        "We define the hyperparameters to be used and train the model.\n",
        "\n",
        "The cell below will print some metrics directly in this notebook, but you may also wish to view the progress of training using TensorBoard. Detailed documentation for using TensorBoard locally can be found [here](https://www.tensorflow.org/guide/summaries_and_tensorboard). Through the Files tab in the sidebar, you can download TensorBoard summary files, which are in `logs/`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4RfQP0M5eUyh"
      },
      "source": [
        "# Feel free to experiment with different values.\n",
        "# A description of all hyperparameters is provided\n",
        "# in the appendix.\n",
        "\n",
        "class BaseHparams(object):\n",
        "  \"\"\"Default hyperparameters.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               total_epochs=100,\n",
        "               learning_rate=0.004,\n",
        "               l2=0.001,\n",
        "               batch_size=256,\n",
        "               window_size=21,\n",
        "               ref_path='hs37d5.fa.gz',\n",
        "               vcf_path='NA12878_calls.vcf.gz',\n",
        "               bam_path='NA12878_sliced.bam',\n",
        "               out_dir='examples',\n",
        "               model_dir='ngs_model',\n",
        "               log_dir='logs'):\n",
        "\n",
        "    self.total_epochs = total_epochs\n",
        "    self.learning_rate = learning_rate\n",
        "    self.l2 = l2\n",
        "    self.batch_size = batch_size\n",
        "    self.window_size = window_size\n",
        "    self.ref_path = ref_path\n",
        "    self.vcf_path = vcf_path\n",
        "    self.bam_path = bam_path\n",
        "    self.out_dir = out_dir\n",
        "    self.model_dir = model_dir\n",
        "    self.log_dir = log_dir"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FuEfufFWshwx"
      },
      "source": [
        "# Delete existing files.\n",
        "!rm -rf examples\n",
        "!rm -rf ngs_model\n",
        "!rm -rf logs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1nDnwaEReile",
        "outputId": "2d9ae5b8-59bc-487c-e8f8-547a8e1543f7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "# This cell should take ~6 minutes to run with the default parameters.\n",
        "hparams = BaseHparams()\n",
        "run(hparams)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Generating data...\n",
            "# of training examples: 7132\n",
            "# of evaluation examples: 2027\n",
            "# of test examples: 1029\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING: Logging before flag parsing goes to stderr.\n",
            "W1007 11:11:44.523533 139916275918720 deprecation.py:323] From <ipython-input-18-e0e9caa56a34>:34: fit_generator (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Please use Model.fit, which supports generators.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Training the model. This should take ~6 minutes...\n",
            "Training complete. Obtaining final metrics...\n",
            "Final eval metrics - loss: 0.072357 - accuracy: 0.996547\n",
            "Final test metrics - loss: 0.070345 - accuracy: 0.997085\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fl4USKB1vUvO"
      },
      "source": [
        "With the default parameters, the final accuracy for this model should be around 99%. Feel free to experiment with different model architectures, learning rates, etc. Though both of the examples we develop are not complex enough to be deployed in production, we hope they will help developers learn to efficiently apply Nucleus and deep learning within genomics.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cEWGUORx63Ou"
      },
      "source": [
        "## Appendix\n",
        "\n",
        "### Hyperparameters\n",
        "\n",
        "Hyperparameter | Description\n",
        "--- | ---\n",
        "`total_epochs` | (int) The number of epochs for which training is run.\n",
        "`learning_rate` | (float) The learning rate for the optimizer.\n",
        "`l2` | (float) The L2 regularization used for the neural network layers.\n",
        "`batch_size` | (int) The number of examples used in one iteration of training, evaluation and testing.\n",
        "`window_size` | (int) The number of bases to consider at once. This should be an odd number so that the middle base is centered evenly.\n",
        "`ref_path` | (str) Path to reference genome.\n",
        "`vcf_path` | (str) Path to truth VCF.\n",
        "`bam_path` | (str) Path to mapped reads.\n",
        "`out_dir` | (str) Path where training, evaluation, and testing TFRecords files written.\n",
        "`model_dir` | (str) Path where model model checkpoints saved. If a checkpoint already exists at this path, training will start from the checkpoint.\n",
        "`log_dir` | (str) Path where TensorBoard logs saved.\n",
        "\n",
        "### Obtaining Original Data Files\n",
        "\n",
        "Below are the commands that were used to obtain and process the data used for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NCy3o7Ew65dN",
        "outputId": "a0e14fe3-4e7a-42f4-f864-7718bad33c0b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 163
        }
      },
      "source": [
        "%%script false\n",
        "# The line above prevents this cell from running as the\n",
        "# default Colab environment does not include the necessary software.\n",
        "\n",
        "# NA12878_sliced.bam\n",
        "samtools view -h \\\n",
        "ftp://ftp-trace.ncbi.nlm.nih.gov/giab/ftp/data/NA12878/NIST_NA12878_HG001_HiSeq_300x/RMNISTHS_30xdownsample.bam \\\n",
        "20:10,000,000-10,100,000 \\\n",
        "-o NA12878_sliced.bam\n",
        "\n",
        "\n",
        "# NA12878_sliced.bam.bai`\n",
        "samtools index NA12878_sliced.bam\n",
        "\n",
        "\n",
        "# NA12878_calls.vcf.gz\n",
        "wget \\\n",
        "ftp://ftp-trace.ncbi.nlm.nih.gov/giab/ftp/release/NA12878_HG001/latest/GRCh37/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer.vcf.gz \\\n",
        "-O NA12878_calls.vcf.gz\n",
        "\n",
        "\n",
        "# NA12878_calls.vcf.gz.tbi\n",
        "wget \\\n",
        "ftp://ftp-trace.ncbi.nlm.nih.gov/giab/ftp/release/NA12878_HG001/latest/GRCh37/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer.vcf.gz.tbi \\\n",
        "-O NA12878_calls.vcf.gz.tbi\n",
        "\n",
        "\n",
        "# hs37d5.fa.gz\n",
        "wget \\\n",
        "ftp://ftp.1000genomes.ebi.ac.uk/vol1/ftp/technical/reference/phase2_reference_assembly_sequence/hs37d5.fa.gz \\\n",
        "-O hs37d5.fa.gz\n",
        "\n",
        "\n",
        "# hs37d5.fa.gz.fai\n",
        "wget \\\n",
        "ftp://ftp.1000genomes.ebi.ac.uk/vol1/ftp/technical/reference/phase2_reference_assembly_sequence/hs37d5.fa.gz.fai \\\n",
        "-O hs37d5.fa.gz.fai\n",
        "\n",
        "\n",
        "# hs37d5.fa.gz.gzi\n",
        "wget \\\n",
        "ftp://ftp.1000genomes.ebi.ac.uk/vol1/ftp/technical/reference/phase2_reference_assembly_sequence/hs37d5.fa.gz.gzi \\\n",
        "-O hs37d5.fa.gz.gzi"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "SyntaxError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;36m  File \u001b[0;32m\"<ipython-input-23-cd78c350cfbe>\"\u001b[0;36m, line \u001b[0;32m6\u001b[0m\n\u001b[0;31m    samtools view -h ftp://ftp-trace.ncbi.nlm.nih.gov/giab/ftp/data/NA12878/NIST_NA12878_HG001_HiSeq_300x/RMNISTHS_30xdownsample.bam 20:10,000,000-10,100,000 -o NA12878_sliced.bam\u001b[0m\n\u001b[0m                ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tsiLCuG3xRQB"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}