gen_chromosome_matrix.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Script for generating chromosome matrix from GRCh37 and dbSNP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "from collections import Counter\n",
    "import pysam\n",
    "import time\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load FASTA & Tabix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_fasta_path = './raw_data/hs37d5.fa' # Reference Fasta File\n",
    "input_tabix_path = './raw_data/GCF_000001405.25.gz' # dbSNP Tabix File\n",
    "output_path = './pretraining_data' # Output Folder Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_fasta(path, header_identifier):\n",
    "    # Read file\n",
    "    lines = open(path,'r').readlines()\n",
    "    lines = list(map(lambda x: x[:-1], lines)) # remove \\n\n",
    "    \n",
    "    # Concatenate sequence data per chromosome\n",
    "    chromosome_dict = {}\n",
    "    chromosome_name = None\n",
    "    chromosome_sequence = ''\n",
    "    for line in lines:\n",
    "        if header_identifier in line:\n",
    "            if chromosome_name:\n",
    "                chromosome_dict[chromosome_name] = ''.join(chromosome_sequence)\n",
    "            # Reinitialize name & sequence\n",
    "            chromosome_name = line[:-1]        \n",
    "            chromosome_sequence = []\n",
    "        else:\n",
    "            # Fill Sequence\n",
    "            chromosome_sequence.append(line)\n",
    "            \n",
    "    # Add the last chromosome\n",
    "    if chromosome_name:\n",
    "        chromosome_dict[chromosome_name] = ''.join(chromosome_sequence)\n",
    "        \n",
    "    return chromosome_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Length GRch37 3137454505\n",
      "Min Length GRch37 4262\n",
      "Max Length GRch37 249250621\n",
      "CPU times: user 19.6 s, sys: 8.04 s, total: 27.6 s\n",
      "Wall time: 29.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "c37_dict = parse_fasta(input_fasta_path,'>')\n",
    "c37_lens = [len(seq) for seq in c37_dict.values()]\n",
    "print('Total Length GRch37', sum(c37_lens))\n",
    "print('Min Length GRch37', min(c37_lens))\n",
    "print('Max Length GRch37', max(c37_lens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb_file = pysam.TabixFile(input_tabix_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Filter only chromosome contigs for FASTA and tabix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fasta_chr_keyword = 'dna:chromosome'\n",
    "fasta_chr_keys = [k for k in c37_dict.keys() if fasta_chr_keyword in k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tb_chr_keyword = 'NC_0000'\n",
    "tabix_chr_keys = [k for k in tb_file.contigs if tb_chr_keyword in k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chr_lens = []\n",
    "for fasta_chr_key in fasta_chr_keys:\n",
    "    chr_lens.append(len(c37_dict[fasta_chr_key]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Construct the chromosome matrix for each chromosome"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_freq(freq_str):\n",
    "    proba_list = []\n",
    "    for src_proba_str in freq_str.split('|'):\n",
    "        probas = src_proba_str.split(':')[-1].split(',')\n",
    "        proba_list.append([0.0 if proba == '.' else float(proba) for proba in probas])\n",
    "        \n",
    "    # Calculate proba\n",
    "    probas = np.array(proba_list).mean(axis=0)    \n",
    "    return probas\n",
    "    \n",
    "def parse_vcf_row(vcf_str):\n",
    "    # parse vcf row\n",
    "    cols = vcf_str.split('\\t')\n",
    "    snv_pos, ref, alts = int(cols[1]) - 1, cols[3], np.array([cols[3]] + cols[4].split(','))\n",
    "    \n",
    "    # retrieve probas from freq info\n",
    "    freq_str = cols[-1].split('FREQ=')[-1].split(';')[0]\n",
    "    probas = parse_freq(freq_str)\n",
    "    \n",
    "    # filter ref_alt and probas\n",
    "    alts, probas = alts[probas > 0], probas[probas > 0]\n",
    "    \n",
    "    return snv_pos, snv_pos + len(ref) - 1, ref, alts, probas\n",
    "\n",
    "# Nucleotide to index (https://en.wikipedia.org/wiki/FASTA_format#Sequence_representation)\n",
    "nt_to_index = { \n",
    "    'A': [0], 'G': [1], 'T': [2], 'C': [3], 'N': [4], 'DEL': [5], 'AI': [6], 'GI': [7], 'TI': [8], 'CI': [9], 'NI': [10],\n",
    "    'U': [2], 'R': [0, 1], 'Y': [2, 3], 'K': [1, 2], 'M': [0, 3], 'S': [1, 3], 'W': [0, 2],\n",
    "    'UI': [8], 'RI': [6, 7], 'YI': [8, 9], 'KI': [7, 8], 'MI': [6, 9], 'SI': [7, 9], 'WI': [6, 8],\n",
    "    'B': [1, 2, 3], 'H': [0, 2, 3], 'V': [0, 1, 3], 'D': [0, 1, 2],\n",
    "    'BI': [7, 8, 9], 'HI': [6, 8, 9], 'VI': [6, 7, 9], 'DI': [6, 7, 8]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2493/2493 [14:45<00:00,  2.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000001.10 | Elapsed time : 1047.1865592002869s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2432/2432 [14:29<00:00,  2.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000002.11 | Elapsed time : 992.5833151340485s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1981/1981 [11:58<00:00,  2.76it/s]\n",
      "  0%|          | 0/1912 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000003.11 | Elapsed time : 798.1739680767059s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1912/1912 [11:32<00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000004.11 | Elapsed time : 794.2839097976685s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1810/1810 [10:50<00:00,  2.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000005.9 | Elapsed time : 758.3790376186371s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1712/1712 [10:23<00:00,  2.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000006.11 | Elapsed time : 720.9959275722504s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1592/1592 [09:49<00:00,  2.70it/s]\n",
      "  0%|          | 0/1464 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000007.13 | Elapsed time : 677.9664719104767s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1464/1464 [08:52<00:00,  2.75it/s]\n",
      "  0%|          | 0/1413 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000008.10 | Elapsed time : 619.2954478263855s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1413/1413 [16:13<00:00,  1.45it/s]\n",
      "  0%|          | 0/1356 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000009.11 | Elapsed time : 1162.5257167816162s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1356/1356 [08:18<00:00,  2.72it/s]\n",
      "  0%|          | 0/1351 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000010.10 | Elapsed time : 570.4666485786438s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1351/1351 [08:07<00:00,  2.77it/s]\n",
      "  0%|          | 0/1339 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000011.9 | Elapsed time : 560.9169895648956s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1339/1339 [08:05<00:00,  2.76it/s]\n",
      "  0%|          | 0/1152 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000012.11 | Elapsed time : 557.5490505695343s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1152/1152 [06:37<00:00,  2.90it/s]\n",
      "  0%|          | 0/1074 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000013.10 | Elapsed time : 469.3514828681946s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1074/1074 [06:11<00:00,  2.89it/s]\n",
      "  0%|          | 0/1026 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000014.8 | Elapsed time : 417.4897701740265s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1026/1026 [05:53<00:00,  2.90it/s]\n",
      "  0%|          | 0/904 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000015.9 | Elapsed time : 394.3665940761566s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 904/904 [05:34<00:00,  2.70it/s]\n",
      "  0%|          | 0/812 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000016.9 | Elapsed time : 369.16913533210754s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 812/812 [04:58<00:00,  2.72it/s]\n",
      "  0%|          | 0/781 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000017.10 | Elapsed time : 333.1558997631073s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 781/781 [04:41<00:00,  2.77it/s]\n",
      "  0%|          | 0/592 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000018.9 | Elapsed time : 324.06333780288696s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 592/592 [03:51<00:00,  2.56it/s]\n",
      "  0%|          | 0/631 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000019.9 | Elapsed time : 256.79016399383545s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 631/631 [03:47<00:00,  2.77it/s]\n",
      "  0%|          | 0/482 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000020.10 | Elapsed time : 255.2924222946167s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 482/482 [02:49<00:00,  2.84it/s]\n",
      "  0%|          | 0/514 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000021.8 | Elapsed time : 190.7699737548828s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 514/514 [02:53<00:00,  2.96it/s]\n",
      "  0%|          | 0/1553 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000022.10 | Elapsed time : 195.44876289367676s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1553/1553 [08:37<00:00,  3.00it/s]\n",
      "  0%|          | 0/594 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000023.10 | Elapsed time : 613.399749994278s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 594/594 [02:49<00:00,  3.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish processing chromosome NC_000024.9 | Elapsed time : 192.52996969223022s\n",
      "CPU times: user 3h 4min 47s, sys: 14min 48s, total: 3h 19min 35s\n",
      "Wall time: 3h 41min 12s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Loop over Chromosome\n",
    "batch_len = 100000\n",
    "for fasta_chr_key, tabix_chr_key, chr_len in zip(fasta_chr_keys, tabix_chr_keys, chr_lens):\n",
    "    start_time = time.time()    \n",
    "    \n",
    "    fasta_seq = c37_dict[fasta_chr_key]  # Fasta reference\n",
    "    data_block = np.zeros((chr_len, 11))  # Target block\n",
    "    \n",
    "    num_batch = (chr_len // batch_len) + 1\n",
    "    last_next_pos = 0\n",
    "    for nb in tqdm.tqdm(range(num_batch)):\n",
    "        start_query_pos, end_query_pos = nb * batch_len, (nb + 1) * batch_len\n",
    "        for row in tb_file.fetch(tabix_chr_key, start_query_pos, end_query_pos):\n",
    "            if ';COMMON' in row:\n",
    "                # parse vcf row\n",
    "                start_pos, end_pos, ref, alts, probas = parse_vcf_row(row)\n",
    "\n",
    "                # handle SNV and INDEL\n",
    "                for alt, proba in zip(alts, probas):\n",
    "                    if len(alt) == len(ref):\n",
    "                        # SNV / MUTATION - assign proba to mutated nucleotide\n",
    "                        for i, nt in enumerate(alt):\n",
    "                            if nt is not 'N':\n",
    "                                data_block[start_pos + i, nt_to_index[nt]] += proba / len(nt_to_index[nt])\n",
    "                    elif len(alt) < len(ref):\n",
    "                        # DELETION\n",
    "                        # assign proba to non-deleted prefix\n",
    "                        for i, nt in enumerate(alt):\n",
    "                            data_block[start_pos + i, nt_to_index[nt]] += proba / len(nt_to_index[nt])\n",
    "\n",
    "                        # assign proba to deleted suffix\n",
    "                        for i in range(len(alt), len(ref)):\n",
    "                            data_block[start_pos + i, nt_to_index['DEL']] += proba / len(nt_to_index['DEL'])\n",
    "                            \n",
    "                    else: # if len(alt) > len(ref):\n",
    "                        # INSERTION\n",
    "                        # assign proba to the prefix nucleotide\n",
    "                        for i in range(len(ref) - 1):\n",
    "                            data_block[start_pos + i, nt_to_index[alt[i]]] += proba / len(nt_to_index[alt[i]])\n",
    "                            \n",
    "                        # assign insertion proba to the last nucleotide\n",
    "                        data_block[start_pos+len(ref)-1, nt_to_index[f'{alt[len(ref)-1]}I']] += proba / len(nt_to_index[f'{alt[len(ref)-1]}I'])\n",
    "\n",
    "                # assign reference to region between SNV\n",
    "                for i in range(last_next_pos, start_pos):\n",
    "                    fasta_nt = fasta_seq[i]\n",
    "                    data_block[i, nt_to_index[fasta_nt]] = 1 / len(nt_to_index[fasta_nt])\n",
    "\n",
    "                # assign new last next position\n",
    "                last_next_pos = start_pos + len(ref)\n",
    "\n",
    "    # assign from last next position onward with nucleotide from reference\n",
    "    for i in range(last_next_pos, chr_len):\n",
    "        fasta_nt = fasta_seq[i]\n",
    "        data_block[i, nt_to_index[fasta_nt]] = 1 / len(nt_to_index[fasta_nt])\n",
    "    \n",
    "    # Normalize over 11D\n",
    "    data_block = data_block / data_block.sum(axis=1, keepdims=True)\n",
    "    \n",
    "    # Dump chromosome data    \n",
    "    np.save(f'{output_path}/{tabix_chr_key}.npy', data_block)\n",
    "    print(f'Finish processing chromosome {tabix_chr_key} | Elapsed time : {time.time() - start_time}s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}