Code-A_pretraining-data-creation-from-HumanRefGenomegene.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import re\n",
    "import pandas as pd\n",
    "from collections import Counter\n",
    "import numpy.random as rand\n",
    "rand.seed(12345)\n",
    "from scipy.stats import bernoulli\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "length_thr= 510 # as mentioned in DNABert paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_gene_data_json(data):\n",
    "    with open('all_chrm_genes.json', 'w') as f:\n",
    "        json.dump(data, f)\n",
    "        \n",
    "def load_json_data(file):\n",
    "    with open(file, 'r') as f:\n",
    "        data = json.loads(f.read())\n",
    "    return data\n",
    "\n",
    "def load_large_json_data(file, num_lines):\n",
    "    with open(file, 'r') as f:\n",
    "        data = json.loads(f.read())\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_base_dir = 'data/'\n",
    "folder = 'GCF_000001405.39' # this file is obtained from the NCBI website https://www.ncbi.nlm.nih.gov/assembly/GCF_000001405.39/ and it contains the Human Reference Genome.\n",
    "\n",
    "files = [f\"chr{id}.fna\" for id in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 'MT', 'X', 'Y']]\n",
    "\n",
    "os.makedirs(os.path.join(data_base_dir, 'cleaned'), exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['chr1.fna', 'chr2.fna', 'chr3.fna', 'chr4.fna', 'chr5.fna', 'chr6.fna', 'chr7.fna', 'chr8.fna', 'chr9.fna', 'chr10.fna', 'chr11.fna', 'chr12.fna', 'chr13.fna', 'chr14.fna', 'chr15.fna', 'chr16.fna', 'chr17.fna', 'chr18.fna', 'chr19.fna', 'chr20.fna', 'chr21.fna', 'chr22.fna', 'chrMT.fna', 'chrX.fna', 'chrY.fna']\n"
     ]
    }
   ],
   "source": [
    "print(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    }
   ],
   "source": [
    "choose= bernoulli.rvs(p=0.5)\n",
    "print(choose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr1.fna  chromosome to be mapped:  chr1.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr2.fna  chromosome to be mapped:  chr2.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr3.fna  chromosome to be mapped:  chr3.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr4.fna  chromosome to be mapped:  chr4.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr5.fna  chromosome to be mapped:  chr5.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr6.fna  chromosome to be mapped:  chr6.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr7.fna  chromosome to be mapped:  chr7.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr8.fna  chromosome to be mapped:  chr8.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr9.fna  chromosome to be mapped:  chr9.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr10.fna  chromosome to be mapped:  chr10.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr11.fna  chromosome to be mapped:  chr11.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr12.fna  chromosome to be mapped:  chr12.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr13.fna  chromosome to be mapped:  chr13.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr14.fna  chromosome to be mapped:  chr14.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr15.fna  chromosome to be mapped:  chr15.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr16.fna  chromosome to be mapped:  chr16.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr17.fna  chromosome to be mapped:  chr17.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr18.fna  chromosome to be mapped:  chr18.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr19.fna  chromosome to be mapped:  chr19.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr20.fna  chromosome to be mapped:  chr20.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr21.fna  chromosome to be mapped:  chr21.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chr22.fna  chromosome to be mapped:  chr22.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chrMT.fna  chromosome to be mapped:  chrMT.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chrX.fna  chromosome to be mapped:  chrX.fna\n",
      "Gene extraction from chromosome file in progress ...\n",
      "File to be mapped:  chrY.fna  chromosome to be mapped:  chrY.fna\n"
     ]
    }
   ],
   "source": [
    "gene_seq_less_thr, chrom_id_list= list(), list()\n",
    "total_genes_cnt= 0\n",
    "\n",
    "for chrom_idx, chrom_id in enumerate(files):\n",
    "    \n",
    "    # extracting the complete sequence length \n",
    "    print('Gene extraction from chromosome file in progress ...')\n",
    "    print('File to be mapped: ', files[chrom_idx], ' chromosome to be mapped: ', chrom_id)\n",
    "    \n",
    "    choose= bernoulli.rvs(p=0.5)\n",
    "    greater_count= 0\n",
    "    \n",
    "    with open(os.path.join(data_base_dir, folder,files[chrom_idx]), \"r\") as f:\n",
    "        lines = f.readlines()\n",
    "        lines = lines[1:]\n",
    "        all_lines= ''.join(lines)\n",
    "        all_lines= re.sub('\\n', '', all_lines)\n",
    "        all_lines=re.sub('N', '', all_lines)\n",
    "        \n",
    "    start_idx= rand.randint(0, min(1000, len(all_lines)))\n",
    "    starting= start_idx\n",
    "    gene_end_idx= len(all_lines) - 1\n",
    "    \n",
    "    while starting < gene_end_idx:        \n",
    "        if bernoulli.rvs(p=0.5) == 0:\n",
    "            gene_len_chosen= length_thr\n",
    "        else:\n",
    "            gene_len_chosen= rand.randint(5,length_thr)\n",
    "\n",
    "        final_seq= all_lines[starting: min((starting + gene_len_chosen+1), gene_end_idx)]\n",
    "        gene_seq_less_thr.append(final_seq)\n",
    "        chrom_id_list.append(chrom_id)\n",
    "\n",
    "        starting= starting + gene_len_chosen + 1  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7641040\n"
     ]
    }
   ],
   "source": [
    "# Only considering genes with length less than 510, results in 14277 genes\n",
    "print(len(gene_seq_less_thr))\n",
    "\n",
    "# with the inclusion of 'N', # of gene sequences is 5.2 million (5201879) only 49 sequences had 'N' in it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({'chr2.fna': 625922, 'chr1.fna': 599559, 'chr3.fna': 515048, 'chr4.fna': 493437, 'chr5.fna': 471516, 'chr6.fna': 442262, 'chr7.fna': 413565, 'chrX.fna': 402908, 'chr8.fna': 377264, 'chr11.fna': 349876, 'chr10.fna': 346282, 'chr12.fna': 346083, 'chr9.fna': 316913, 'chr13.fna': 254954, 'chr14.fna': 235561, 'chr15.fna': 220047, 'chr17.fna': 215798, 'chr16.fna': 212684, 'chr18.fna': 208094, 'chr20.fna': 166363, 'chr19.fna': 152041, 'chr21.fna': 104177, 'chr22.fna': 101872, 'chrY.fna': 68772, 'chrMT.fna': 42})\n"
     ]
    }
   ],
   "source": [
    "print(Counter(chrom_id_list))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting gene sequence to k-mer form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Taken from DNAbert's codebase https://github.com/jerryji1993/DNABERT - motif/motif_utils.py\n",
    "def seq2kmer(seq, k):\n",
    "    \"\"\"\n",
    "    Convert original sequence to kmers\n",
    "    \n",
    "    Arguments:\n",
    "    seq -- str, original sequence.\n",
    "    k -- int, kmer of length k specified.\n",
    "    \n",
    "    Returns:\n",
    "    kmers -- str, kmers separated by space\n",
    "\n",
    "    \"\"\"\n",
    "    kmer = [seq[x:x+k] for x in range(len(seq)+1-k)]\n",
    "    kmers = \" \".join(kmer)\n",
    "    return kmers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9377\n"
     ]
    }
   ],
   "source": [
    "# remove duplicates from a list\n",
    "gene_seq_less_thr_set = list(set(gene_seq_less_thr))\n",
    "print(len(gene_seq_less_thr)- len(gene_seq_less_thr_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7631663\n"
     ]
    }
   ],
   "source": [
    "# pretrain_seq\n",
    "print(len(gene_seq_less_thr_set))\n",
    "\n",
    "# divide the data into 90% training, 10% validation data\n",
    "choose_seq= bernoulli.rvs(size= len(gene_seq_less_thr_set), p=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "       0, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "       1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "       1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1,\n",
       "       1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "choose_seq[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_pre_data= open('dnabert_original_train_pretraining_data.txt', 'w')\n",
    "valid_pre_data= open('dnabert_original_valid_pretraining_data.txt', 'w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "for gene_seq_id, gene_seq in enumerate(gene_seq_less_thr_set):\n",
    "    kmer_form= seq2kmer(gene_seq, 6)\n",
    "    \n",
    "    if choose_seq[gene_seq_id] == 1: # use it for training\n",
    "        train_pre_data.write(kmer_form + '\\n')\n",
    "    else: # use it for evaluation\n",
    "        valid_pre_data.write(kmer_form + '\\n')   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_pre_data.close()\n",
    "valid_pre_data.close()"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}