hdr-modeller.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import re\n",
    "\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.utils import compute_class_weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "LENGTH_OF_PAM = 3\n",
    "rev = {'A':'T', 'T':'A', 'C':'G', 'G':'C', 'a':'t', 't':'a', 'c':'g', 'g':'c'}\n",
    "nucleotides = ['A', 'T', 'C', 'G']\n",
    "dinucleotides = [i + j for i in nucleotides for j in nucleotides]\n",
    "allnucleotides = nucleotides + dinucleotides\n",
    "TOKENS_SI = [(1, i) for i in nucleotides]\n",
    "TOKENS_DI = [(2, i) for i in dinucleotides]\n",
    "TOKENS = TOKENS_SI + TOKENS_DI\n",
    "\n",
    "CLASS_LABEL = 'hdr_0.2'\n",
    "RANDOM_STATE = 3463767\n",
    "N_TREES = 2000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(file_name):\n",
    "    df = pd.read_csv(file_name)\n",
    "    print(f'Read {len(df)} targets from \"{file_name}\"')\n",
    "    return df\n",
    "\n",
    "\n",
    "## Remove owasites with below the specified cleavage threshold\n",
    "def remove_unwanted(df, threshold):\n",
    "    return df[ (df.indels / df.trials) > threshold]\n",
    "\n",
    "\n",
    "## Returns the reverse complement to a DNA sequence\n",
    "def reverse_complement(sequence):\n",
    "    return ''.join([rev[i] for i in sequence[::-1]])\n",
    "\n",
    "\n",
    "## Breaks an oligo down into it's arms (at the point mutation)\n",
    "def get_oligo_components(guide, oligo, distance_from_pam, mutation_size, strand,\n",
    "                         truncation=0):\n",
    "    \n",
    "    pam_mutation_index = range(20 - distance_from_pam - mutation_size,\n",
    "                               20 - distance_from_pam)\n",
    "    \n",
    "    guide_regex = ''.join([nucleotide if i not in pam_mutation_index\n",
    "                           else '[ATCG]' for i, nucleotide in enumerate(guide)])\n",
    "    match = re.search(guide_regex, oligo, flags=re.IGNORECASE)\n",
    "    if match:\n",
    "        dist = match.span()[1] - distance_from_pam - 3\n",
    "        arms = [oligo[: dist-1].upper(), oligo[dist:].upper()]\n",
    "    else:\n",
    "        oligo = reverse_complement(oligo)\n",
    "        match = re.search(guide_regex, oligo, flags=re.IGNORECASE)\n",
    "        dist = match.span()[1] - distance_from_pam - 3\n",
    "        arms = [oligo[: dist-1].upper(), oligo[dist:].upper()]\n",
    "    if truncation > 0:\n",
    "        oligo_arm1 = arms[0][-truncation:]\n",
    "        oligo_arm2 = arms[1][0: truncation]\n",
    "    else:\n",
    "        oligo_arm1 = arms[0]\n",
    "        oligo_arm2 = arms[1]\n",
    "    return oligo_arm1, oligo_arm2, oligo_arm1 + oligo_arm2\n",
    "\n",
    "\n",
    "## Get arms from oligo, optionally with a maximum length\n",
    "def get_truncated_arms(df, truncation):\n",
    "    return df.apply(\n",
    "        lambda _: pd.Series(\n",
    "            get_oligo_components(_.full_guide_sequence, _.plus_oligo,\n",
    "                                 _.distance_from_pam, _.mutation_size,\n",
    "                                 _.strand, truncation),\n",
    "            index=['oligo_arm1', 'oligo_arm2', 'oligo_trimmed']), axis=1\n",
    "    )\n",
    "\n",
    "\n",
    "## Ensure all items are in the same orientation (+)\n",
    "def get_plus_oligo(guide, oligo, distance_from_pam, mutation_size):\n",
    "    spacer_size = 20\n",
    "    pam_mutation_index = range(\n",
    "        spacer_size - distance_from_pam - mutation_size,\n",
    "        spacer_size - distance_from_pam\n",
    "    )\n",
    "    guide_regex = ''.join([nucleotide if i not in pam_mutation_index\n",
    "                           else '[ATCG]' for i,\n",
    "                           nucleotide in enumerate(guide)])\n",
    "    match = re.search(guide_regex, oligo, flags=re.IGNORECASE)\n",
    "    if match:\n",
    "        return oligo\n",
    "    oligo_rev = reverse_complement(oligo)\n",
    "    match = re.search(guide_regex, oligo_rev, flags=re.IGNORECASE)\n",
    "    if match:\n",
    "        return oligo_rev        \n",
    "\n",
    "\n",
    "## Merge experiments using the same oligo\n",
    "def merge_duplicates(df, merge_key):\n",
    "    tmp_df = df.reset_index().groupby(merge_key, group_keys=False).agg(\n",
    "        {'index': 'first',\n",
    "         'point_mutations':sum,\n",
    "         'indels': sum,\n",
    "         'trials': sum}).set_index('index')\n",
    "    cols_to_use = df.columns.difference(tmp_df.columns)\n",
    "\n",
    "    merged_df = pd.concat([df[cols_to_use], tmp_df], axis=1, join='inner')\n",
    "    return merged_df\n",
    "\n",
    "\n",
    "## Generate binary labels\n",
    "def process_labels(df):\n",
    "    labels_df = df.loc[:, ['indels', 'point_mutations', 'trials']]\n",
    "    labels_df.loc[:, 'nhej_ratio'] = (labels_df.indels / labels_df.trials)\n",
    "    for i in [0.4, 0.5, 0.6]:\n",
    "        labels_df.loc[:, 'nhej_{}'.format(i)] = \\\n",
    "            (labels_df.indels / labels_df.trials) > i\n",
    "    labels_df.loc[:, 'hdr_ratio'] = (labels_df.point_mutations / labels_df.trials)\n",
    "    for i in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]:\n",
    "        labels_df.loc[:, 'hdr_{}'.format(i)] = \\\n",
    "            (labels_df.point_mutations / labels_df.trials) > i\n",
    "    labels_df.loc[:, 'hdr_all'] = \\\n",
    "        (labels_df.point_mutations / (labels_df.point_mutations + labels_df.trials))\n",
    "    for i in [0.1, 0.18, 0.2, 0.3, 0.4, 0.5, 0.6]:\n",
    "        labels_df.loc[:, 'hdr_all_{}'.format(i)] = (\n",
    "            labels_df.point_mutations /\n",
    "            (labels_df.point_mutations + labels_df.trials)\n",
    "        ) > i\n",
    "    return labels_df\n",
    "\n",
    "\n",
    "## Features from guide sequence\n",
    "def process_features_guide_spacer(df):   \n",
    "    prefixes = {'si': 1, 'di': 2}\n",
    "    new_df = pd.DataFrame(data=None, index=df.index)\n",
    "    for pre in prefixes:\n",
    "        l = prefixes[pre] - 1\n",
    "        for i in range(1,21 - l):\n",
    "            new_df.loc[:, 'guide_{}_{:02d}'.format(pre, i)] = \\\n",
    "                df.full_guide_sequence.str[\n",
    "                    - LENGTH_OF_PAM - i - l: - LENGTH_OF_PAM - i + 1\n",
    "                ]\n",
    "    global_guide_df = new_df.apply(pd.Series.value_counts, axis=1) \\\n",
    "        .reindex(allnucleotides, axis=1).fillna(0).astype(np.int8) / 20 * 100\n",
    "    return global_guide_df\n",
    "\n",
    "\n",
    "def process_features(df, truncation = 0):\n",
    "    df.loc[:, 'plus_oligo'] = np.vectorize(get_plus_oligo) \\\n",
    "        (df.full_guide_sequence, df.ss_oligo, df.distance_from_pam, df.mutation_size)\n",
    "\n",
    "    guide_df = process_features_guide_spacer(df)\n",
    "    \n",
    "    ## Create df with the oligo components (arms)\n",
    "    oligo_arms_df = get_truncated_arms(df, truncation)\n",
    "    ## Initialize empty df for oligo nucleotide count\n",
    "    oligo_df = pd.DataFrame(data=None, index=df.index)\n",
    "    ## Add nucleotide counts to above df\n",
    "    for i in TOKENS:\n",
    "        oligo_df.loc[:, 'oligo_arm2_{}'.format(i[1])] = \\\n",
    "            oligo_arms_df.oligo_arm2.apply(\n",
    "                lambda x: (\n",
    "                    len(\n",
    "                        re.findall('(?={})'.format(i[1]), x, flags=re.IGNORECASE)\n",
    "                    ) / (len(x)/i[0]) * 100\n",
    "                )\n",
    "            )\n",
    "\n",
    "    return pd.concat([guide_df, oligo_df], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train random forest model\n",
    "Read from a CSV file containing these columns:\n",
    "- point_mutations\n",
    "- indels\n",
    "- trials\n",
    "\n",
    "\n",
    "- mutation_size\n",
    "- distance_from_pam\n",
    "- full_guide_sequence\n",
    "- ss_oligo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read 98 targets from \"data/train_set.csv\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RandomForestClassifier(bootstrap=True,\n",
       "            class_weight={False: 1.28125, True: 0.82}, criterion='gini',\n",
       "            max_depth=None, max_features='auto', max_leaf_nodes=None,\n",
       "            min_impurity_decrease=0.0, min_impurity_split=None,\n",
       "            min_samples_leaf=1, min_samples_split=2,\n",
       "            min_weight_fraction_leaf=0.0, n_estimators=2000, n_jobs=1,\n",
       "            oob_score=True, random_state=3463767, verbose=0,\n",
       "            warm_start=False)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file_name = 'data/train_set.csv'\n",
    "\n",
    "df = read_data(file_name)\n",
    "df = merge_duplicates(df, 'ss_oligo')\n",
    "df = remove_unwanted(df, 0.4)\n",
    "\n",
    "X = process_features(df)\n",
    "y = process_labels(df)['hdr_0.2']\n",
    "\n",
    "class_weights = dict(\n",
    "    zip(\n",
    "        np.unique(y),\n",
    "        compute_class_weight('balanced', np.unique(y), y)\n",
    "    )\n",
    ")\n",
    "\n",
    "forest = RandomForestClassifier(\n",
    "    n_estimators=N_TREES,\n",
    "    oob_score=True,\n",
    "    class_weight = class_weights,\n",
    "    random_state=RANDOM_STATE\n",
    ")\n",
    "\n",
    "forest.fit(X, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classify targets\n",
    "\n",
    "Read from a CSV file containing these columns:\n",
    "- mutation_size\n",
    "- distance_from_pam\n",
    "- full_guide_sequence\n",
    "- ss_oligo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read 6 targets from \"data/test_oligos.txt\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([ True,  True,  True, False,  True, False])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file_name = 'data/test_oligos.txt'\n",
    "\n",
    "df = read_data(file_name)\n",
    "X = process_features(df)\n",
    "forest.predict(X)"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}