Codon Optimization Dataset Creation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Codon Optimization Dataset Creation\n",
    "\n",
    "To begin, let's create a 50/50 class split of optimized and non-optimized of sequences. To do this, we'll load up all our RefSeq CDSs and then randomly sort them into either class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "version = \"0.1.3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(201179, 201179)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from random import shuffle, seed, choice\n",
    "from Bio import SeqIO\n",
    "seed(42) # for reproducibility\n",
    "\n",
    "# load the sequences and shuffle them randomly\n",
    "CDS = list(SeqIO.parse(\"inserts/all_refseq_prokaryote_CDS.fasta\", \"fasta\"))\n",
    "CDS = list(set([str(x.seq).upper() for x in CDS if len(x.seq) % 3 == 0])) # to remove duplicates and incomplete CDSs\n",
    "CDS = [seq for seq in CDS if set(seq) == {\"A\", \"T\", \"G\", \"C\"}] # restrict to only ATGC characters\n",
    "shuffle(CDS)\n",
    "\n",
    "# split them 50/50 and remove the first codon, which will always translate to M\n",
    "optimized, natural = CDS[:int(len(CDS)/2)], CDS[int(len(CDS)/2):]\n",
    "optimized, natural = [seq[3:] for seq in optimized], [seq[3:] for seq in natural] # remove the first codon, which is always translated as an M\n",
    "len(optimized), len(natural)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Target Data Collection\n",
    "\n",
    "Before we can do our optimization, we'll need to collect the data on our targets, which will be highly expressed and overall genome CUB.\n",
    "\n",
    "### Highly Expressed Genes (HEGs) Targets\n",
    "The HEGs are stored as a list of FASTA files in the targets/heg directory. We are going to iterate over them and calculate their codon usage bias."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import freqgen\n",
    "\n",
    "heg_cub = []\n",
    "heg_names = []\n",
    "\n",
    "for heg in os.listdir(\"targets/heg\"):\n",
    "    \n",
    "    # only deal with fasta files\n",
    "    if not heg.endswith(\".fasta\"):\n",
    "        continue\n",
    "        \n",
    "    # calculate the HEGs' CUB\n",
    "    seq = [x.seq for x in SeqIO.parse(f'targets/heg/{heg}', \"fasta\")]\n",
    "    try:\n",
    "        heg_cub.append(freqgen.codon_frequencies(seq))\n",
    "        heg_names.append(heg)\n",
    "    except ValueError: # ignore bad codons\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Genome Targets\n",
    "In addition to targeting various species' HEGs, we will also target overall genome CUB. To do so, we are using the data from [this paper](https://doi.org/10.1186/s12859-017-1793-7). It is stored as a large CSV file, so we'll use pandas to read it in and convert it to list of dicts of CUB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/BenjaminLee/Desktop/Python/Research/condo/env/lib/python3.6/importlib/_bootstrap.py:205: RuntimeWarning: numpy.dtype size changed, may indicate binary incompatibility. Expected 96, got 88\n",
      "  return f(*args, **kwds)\n",
      "/Users/BenjaminLee/Desktop/Python/Research/condo/env/lib/python3.6/site-packages/IPython/core/interactiveshell.py:2785: DtypeWarning: Columns (5) have mixed types. Specify dtype option on import or set low_memory=False.\n",
      "  interactivity=interactivity, compiler=compiler, result=result)\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv(\"targets/genome/o563612-refseq_species.tsv\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 114111/114111 [02:45<00:00, 688.61it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "\n",
    "genome_cub = [] # a list of the each species' CUB\n",
    "genome_names = [] # the name of each species\n",
    "\n",
    "for i in trange(len(df)):\n",
    "    genome_cub.append(dict(df.iloc[i][12:])) # the raw codon counts\n",
    "    genome_names.append(df.iloc[i][3]) # species name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "114307"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# how many targets are there?\n",
    "targets = genome_cub + heg_cub\n",
    "len(targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimize the sequences\n",
    "\n",
    "Ok, now we have all our targets and inserts. We'll iterate over the CDSs we've selected for optimization and optimize them using one of the methods towards one of the targets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Bio.Data.CodonTable as ct\n",
    "\n",
    "def _synonymous_codons(genetic_code_dict):\n",
    "\n",
    "    # invert the genetic code dictionary to map each amino acid to its codons\n",
    "    codons_for_amino_acid = {}\n",
    "    for codon, amino_acid in genetic_code_dict.items():\n",
    "        codons_for_amino_acid[amino_acid] = codons_for_amino_acid.get(amino_acid, [])\n",
    "        codons_for_amino_acid[amino_acid].append(codon)\n",
    "\n",
    "    # create dictionary of synonymous codons\n",
    "    # Example: {'CTT': ['CTT', 'CTG', 'CTA', 'CTC', 'TTA', 'TTG'], 'ATG': ['ATG']...}\n",
    "    return {codon: codons_for_amino_acid[genetic_code_dict[codon]] for codon in genetic_code_dict.keys()}\n",
    "\n",
    "\n",
    "synonymous_codons = {k: _synonymous_codons(v.forward_table) for k, v in ct.unambiguous_dna_by_id.items()}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def absolute_to_relative(target, trans_table):\n",
    "    relative = {}\n",
    "    for i in synonymous_codons[trans_table].keys():\n",
    "        try:\n",
    "            relative[i] = target[i] / sum((target[codon] for codon in synonymous_codons[trans_table][i]))\n",
    "        except ZeroDivisionError:\n",
    "            relative[i] = 1 / len(synonymous_codons[trans_table][i]) # if an amino acid is never used in the reference set, then all its codons are used equally\n",
    "    return relative"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Bio.Seq import Seq\n",
    "import numpy as np\n",
    "np.random.seed(42)\n",
    "\n",
    "def CAI_max(seq, absolute, trans_table):\n",
    "    '''The CAI maximization method, which takes the most used codon for each AA'''\n",
    "    relative = absolute_to_relative(absolute, trans_table)\n",
    "    optimized_seq = \"\"\n",
    "    for codon in [seq[i:i + 3] for i in range(0, len(seq), 3)]:\n",
    "        optimized_seq += max([(c, relative[c]) for c in synonymous_codons[trans_table][codon]], key=lambda x: x[1])[0] # chose the most used synonym\n",
    "    return optimized_seq\n",
    "\n",
    "def multinomial(seq, absolute, trans_table):\n",
    "    '''The multinomial method, which chooses codons with repect to their usage frequencies'''\n",
    "    relative = absolute_to_relative(absolute, trans_table)\n",
    "    optimized_seq = \"\"\n",
    "    for codon in [seq[i:i + 3] for i in range(0, len(seq), 3)]:\n",
    "        optimized_seq += np.random.choice(synonymous_codons[trans_table][codon], p=[relative[c] for c in synonymous_codons[trans_table][codon]]) # pick a codon with the frequencies of the target\n",
    "    return optimized_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize(seq):\n",
    "    # choose a codon optimization method\n",
    "    target_method = np.random.choice([\"cai_max\", \"multinomial\"]) \n",
    "    \n",
    "    # choose a target\n",
    "    targets = [heg_cub, genome_cub]\n",
    "    target_type = np.random.choice(targets)\n",
    "    target_idx = np.random.randint(len(target_type))\n",
    "    target = target_type[target_idx]\n",
    "    target_name = [heg_names, genome_names][targets.index(target_type)][target_idx]\n",
    "    \n",
    "    # use an alternate translation table, if known\n",
    "    if target_type == genome_cub:\n",
    "        trans_table = df.iloc[target_idx][\"Translation Table\"]\n",
    "    else:\n",
    "        trans_table = 11\n",
    "        \n",
    "    # perform the optimization\n",
    "    if target_method == \"cai_max\":\n",
    "        result = CAI_max(seq, target, trans_table)\n",
    "    elif target_method == \"multinomial\":\n",
    "        result = multinomial(seq, target, trans_table)\n",
    "    \n",
    "    # return a results dict\n",
    "    return dict(sequence=result,\n",
    "                method=target_method, \n",
    "                target_type=[\"heg\", \"genome\"][targets.index(target_type)], \n",
    "                target_name=target_name,\n",
    "                trans_table=trans_table,\n",
    "                optimized=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 201179/201179 [47:52<00:00, 70.05it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "results = []\n",
    "for sequence in tqdm(optimized):\n",
    "    try:\n",
    "        results.append(optimize(sequence))\n",
    "    except KeyError:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output a HDF5 file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 201179/201179 [00:00<00:00, 272684.17it/s]\n"
     ]
    }
   ],
   "source": [
    "for sequence in tqdm(natural):\n",
    "    results.append(dict(sequence=sequence,\n",
    "                        optimized=0,\n",
    "                        method=None,\n",
    "                        target_type=None,\n",
    "                        target_name=None,\n",
    "                        trans_table=11,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = [\"sequence\", \"optimized\", \"method\", \"trans_table\", \"target_type\", \"target_name\"]\n",
    "pd.DataFrame(results, columns=columns).to_hdf(f\"data/condo-{version}.h5\", \"condo\")"
   ]
  }
 ],
 "metadata": {
  "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.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}