MemMaps Production Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "de26e100",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "from itertools import groupby\n",
    "import itertools\n",
    "from operator import itemgetter\n",
    "import os\n",
    "import pyBigWig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af4bdc3",
   "metadata": {},
   "source": [
    "First, we need to define the location of the fasta file, the gap file, the directory with the targets, and the directories, where to save the input and target memory maps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cf9c4dba",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets_dir = '/data/users/goodarzilab/darya/work/Datasets'\n",
    "# target_dir = '/data/users/goodarzilab/darya/work/Datasets/corderoLabData/bwOS.PDX.ATAC'\n",
    "target_dir = '/data/users/goodarzilab/darya/work/Datasets'\n",
    "\n",
    "memmap_data_dir = os.path.join(os.getcwd(), 'hg_19_mmaps_upd')\n",
    "targets_memmap_data_dir = os.path.join(os.getcwd(), 'targets_memmaps_test')\n",
    "\n",
    "\n",
    "fasta_file = os.path.join(datasets_dir, 'hg19.ml.fa')\n",
    "gaps_file = os.path.join(datasets_dir, 'hg19_gaps.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e047bf61",
   "metadata": {},
   "source": [
    "Then, we define the list with full directories of each of the targets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9b33fbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "targets_list = os.listdir(target_dir)\n",
    "targets_path_list = [os.path.join(target_dir, tar) for tar in targets_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f36e932d",
   "metadata": {},
   "source": [
    "Then, we need to create a dictionary with nucleotides - both lower and upper case. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53e17e57",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['A', 'T', 'C', 'G', 'N', 'a', 't', 'c', 'g', 'n'], dtype='<U1')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nucs = np.array([\"A\", \"T\", \"C\", \"G\", \"N\"])\n",
    "mapping = {u:i for i,u in enumerate(nucs)}\n",
    "nucs_lower = [nuc.lower() for nuc in nucs]\n",
    "mapping_lower = {u:i for i,u in enumerate(nucs_lower)}\n",
    "\n",
    "mapping.update(mapping_lower)\n",
    "nucs = np.append(nucs, nucs_lower)\n",
    "nucs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b184cbd5",
   "metadata": {},
   "source": [
    "Next, we open the gaps file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c3448d5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = []\n",
    "with open(gaps_file) as f:\n",
    "    for line in f:\n",
    "        lines.append(line.split())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50ea0dbc",
   "metadata": {},
   "source": [
    "For the sake of having a full list of chromosome names, we pul it from a random bw file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "03c57db3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'chr1': 249250621,\n",
       " 'chr10': 135534747,\n",
       " 'chr11': 135006516,\n",
       " 'chr12': 133851895,\n",
       " 'chr13': 115169878,\n",
       " 'chr14': 107349540,\n",
       " 'chr15': 102531392,\n",
       " 'chr16': 90354753,\n",
       " 'chr17': 81195210,\n",
       " 'chr18': 78077248,\n",
       " 'chr19': 59128983,\n",
       " 'chr2': 243199373,\n",
       " 'chr20': 63025520,\n",
       " 'chr21': 48129895,\n",
       " 'chr22': 51304566,\n",
       " 'chr3': 198022430,\n",
       " 'chr4': 191154276,\n",
       " 'chr5': 180915260,\n",
       " 'chr6': 171115067,\n",
       " 'chr7': 159138663,\n",
       " 'chr8': 146364022,\n",
       " 'chr9': 141213431,\n",
       " 'chrX': 155270560}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# target = pyBigWig.open(os.path.join(target_dir, 'PDX008_1.bw'), 'r')\n",
    "target = pyBigWig.open(os.path.join(target_dir, 'CNhs12843.bw'), 'r')\n",
    "\n",
    "chrom_seq = target.chroms() \n",
    "chrom_seq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e798e90b",
   "metadata": {},
   "source": [
    "With the gap information derived above, we create a dictionary where each cromosome name corresponds to the \"non-gap\" locations, i.e., with the gaps extracted. \n",
    "Different .bed files may be differently organized, so you need to manually define the following variables: chrom_position , gaps_start, gaps_end that correspond to the location of the chromosome name and the boundaries of the gap. \n",
    "See example below. \n",
    "Note that you need to have chrom_seq values match the values from both inputs and targets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "99f3686b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['0',\n",
       " 'chr1',\n",
       " '124535434',\n",
       " '142535434',\n",
       " '1271',\n",
       " 'N',\n",
       " '18000000',\n",
       " 'heterochromatin',\n",
       " 'no']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lines[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9732f5fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "chrom_position = 1\n",
    "gaps_start = 2\n",
    "gaps_end = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cf85522a",
   "metadata": {},
   "outputs": [],
   "source": [
    "chrom_gaps = dict()\n",
    "for i, g in itertools.groupby(lines, lambda x: x[chrom_position]):\n",
    "    if i in list(chrom_seq.keys()):\n",
    "        arr =  [np.array(el[gaps_start:gaps_end + 1]).astype(int) for el in g]\n",
    "        arr = sorted(arr, key=itemgetter(0))\n",
    "        arr_corrected = [[int(arr[i][1]), int(arr[i+1][0])] for i in range(len(arr)-1) if int(arr[i][1]) != int(arr[i+1][0])]\n",
    "        if int(arr[-1][1]) != chrom_seq[i]:\n",
    "            arr_corrected.append([int(arr[-1][1]), chrom_seq[i]])        \n",
    "        chrom_gaps[i] = arr_corrected"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ed3c31",
   "metadata": {},
   "source": [
    "The iterator below opens a fast file, and for each header (chromosome), extracts the fill sequence as a long string. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e04828fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fasta_iter(fasta_name):\n",
    "    \"\"\"\n",
    "    modified from Brent Pedersen\n",
    "    Correct Way To Parse A Fasta File In Python\n",
    "    given a fasta file. yield tuples of header, sequence\n",
    "    \"\"\"\n",
    "    fh = open(fasta_name)\n",
    "    faiter = (x[1] for x in groupby(fh, lambda line: line[0] == \">\"))\n",
    "\n",
    "    for header in faiter:\n",
    "        headerStr = header.__next__()[1:].strip()\n",
    "        seq = \"\".join(s.strip() for s in faiter.__next__())\n",
    "\n",
    "        yield (headerStr, seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70916e2e",
   "metadata": {},
   "source": [
    "The loop below uses the iterator defined above, opens the fasta file, extracts the gaps, and saves the result as a memory map. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b48de0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fiter = fasta_iter(fasta_file)\n",
    "for ff in fiter:\n",
    "    headerStr, seq = ff\n",
    "    seq = np.array(list(seq))\n",
    "    print (headerStr)\n",
    "    for key in chrom_gaps.keys():\n",
    "        if key == headerStr:\n",
    "            print (key)\n",
    "            memmap_name = key + \"_sequence.dta\"\n",
    "            els_corrected = chrom_gaps[key]\n",
    "#             lengths.insert(0, 0)\n",
    "            lengths = [e[1] - e[0] for e in els_corrected]\n",
    "            contigs_shape = sum(lengths)\n",
    "            lengths.insert(0, 0)\n",
    "            lengths_idx = np.cumsum(lengths)\n",
    "            chr_contigs =  np.memmap(os.path.join(memmap_data_dir, (key + \"_contigs_sequence.dta\")), mode='w+', dtype='float32', shape=(contigs_shape))\n",
    "            for i in range(len(els_corrected)):\n",
    "#                 print (i)\n",
    "                seq_subset = seq[els_corrected[i][0]:els_corrected[i][1]]\n",
    "                seq_subset_mapped = np.zeros((len(seq_subset),))\n",
    "                for nuc in nucs:\n",
    "                    j = np.where(seq_subset[0:len(seq_subset)] == nuc)[0]\n",
    "                    seq_subset_mapped[j] = mapping[nuc]\n",
    "#                 print (seq_subset_mapped.shape, lengths_idx[i+1] - lengths_idx[i])\n",
    "                chr_contigs[lengths_idx[i]:lengths_idx[i+1]] = seq_subset_mapped\n",
    "            chr_contigs.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b05917d",
   "metadata": {},
   "source": [
    "The loop below uses the iterator defined above, opens the targets bw files, extracts the gaps, and saves the result as a memory map, where all the gap-free targets are sequentially stacked. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d384cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "targets_path_list = [os.path.join(target_dir, tar) for tar in targets_list]\n",
    "for key in chrom_gaps.keys():\n",
    "    print (key)\n",
    "    memmap_name = key + \"_sequence.dta\"\n",
    "    els_corrected = chrom_gaps[key]\n",
    "    lengths = [e[1] - e[0] for e in els_corrected]\n",
    "    contigs_shape = sum(lengths)\n",
    "    lengths.insert(0, 0)\n",
    "    lengths_idx = np.cumsum(lengths)\n",
    "    chr_contigs =  np.memmap(os.path.join(targets_memmap_data_dir, (key + \"_targets_contigs_sequence.dta\")), mode='w+', dtype='float32', shape=(len(targets_list), contigs_shape))\n",
    "    for i in range(len(els_corrected)):\n",
    "        for j in range(len(targets_path_list)):\n",
    "            target = pyBigWig.open(targets_path_list[j], 'r')\n",
    "            chr_contigs[j, lengths_idx[i]:lengths_idx[i+1]] = target.values(key, els_corrected[i][0], els_corrected[i][1])\n",
    "    chr_contigs.flush()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "basenji_pytorch",
   "language": "python",
   "name": "basenji_pytorch"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}