Bug fixes

Add unit tests and fix the bugs that cause unexpected bahaviors during test

Author: yztxwd

seqchromloader/__init__.py

@@ -1,3 +1,3 @@
 from .loader import SeqChromDatasetByDataFrame, SeqChromDatasetByBed, SeqChromDatasetByWds, SeqChromDataModule
 from .writer import dump_data_webdataset, convert_data_webdataset
-from .utils import filter_chromosomes, make_random_shift, make_flank, chop_genome, dna2OneHot, rev_comp
+from .utils import filter_chromosomes, make_random_shift, make_flank, chop_genome, dna2OneHot, rev_comp, get_genome_sizes, random_coords

seqchromloader/loader.py

@@ -8,7 +8,7 @@
 import torch
 import random
 import pysam
-import pyfasta
+import pyfaidx
 import pyBigWig
 import numpy as np
 import pandas as pd
@@ -120,7 +120,7 @@ def __init__(self,
 
         self.dataframe = dataframe        
         self.genome_fasta = genome_fasta
-        self.genome_pyfasta = None
+        self.genome_pyfaidx = None
         self.bigwig_filelist = bigwig_filelist
         self.bigwigs = None
         self.target_bam = target_bam
@@ -133,7 +133,7 @@ def __init__(self,
     def initialize(self):
         # create the stream handler after child processes spawned to enable parallel reading
         # this function will be called by worker_init_function in DataLoader
-        self.genome_pyfasta = pyfasta.Fasta(self.genome_fasta)
+        self.genome_pyfaidx = pyfaidx.Fasta(self.genome_fasta)
         self.bigwigs = [pyBigWig.open(bw) for bw in self.bigwig_filelist]
         if self.target_bam is not None:
             self.target_pysam = pysam.AlignmentFile(self.target_bam)
@@ -149,7 +149,7 @@ def __getitem__(self, idx):
                 item.start,
                 item.end,
                 item.label,
-                genome_pyfasta=self.genome_pyfasta,
+                genome_pyfaidx=self.genome_pyfaidx,
                 bigwigs=self.bigwigs,
                 target_bam=self.target_pysam,
                 strand=item.strand,

seqchromloader/utils.py

@@ -28,7 +28,9 @@ def get_genome_sizes(gs=None, genome=None, to_filter=None, to_keep=None):
     elif genome:
         genome_sizes = (pd.DataFrame(chromsizes(genome))
                         .T
-                        .rename(columns={0:"chrom", 1:"len"}))
+                        .reset_index()
+                        .rename(columns={"index":"chrom", 0:"start", 1:"end"})
+                        .assign(length=lambda x: x["end"] - x["start"]))[["chrom", "length"]]
     else:
         raise Exception("Either gs or genome should be provided!")
 
@@ -74,7 +76,7 @@ def filter_chromosomes(coords, to_filter=None, to_keep=None):
         corods_out = coords
     return corods_out
 
-def make_random_shift(coords, L, buffer=25):
+def make_random_shift(coords, L, buffer=0):
     """
     This function takes as input a set of bed coordinates dataframe 
     It finds the mid-point for each record or Interval in the bed file,
@@ -152,8 +154,8 @@ def random_coords(gs:str=None, genome:str=None, incl:BedTool=None, excl:BedTool=
     else:
         raise Exception("Either gs or genome should be provided!")
 
-    if incl: shuffle_kwargs.update({"incl": incl})
-    if excl: shuffle_kwargs.update({"excl": excl})
+    if incl: shuffle_kwargs.update({"incl": incl.fn})
+    if excl: shuffle_kwargs.update({"excl": excl.fn})
 
     return (BedTool()
             .random(l=l, n=n, **random_kwargs)
@@ -194,7 +196,7 @@ def intervals_loop(chrom, start, stride, l, size):
 
     genome_sizes = get_genome_sizes(gs=gs, genome=genome, to_keep=chroms)
 
-    genome_chops = pd.concat([intervals_loop(i.Index, 0, stride, l, i.len) 
+    genome_chops = pd.concat([intervals_loop(i.chrom, 0, stride, l, i.length) 
                                 for i in genome_sizes.itertuples()])
     genome_chops_bdt = BedTool.from_dataframe(genome_chops)
 
@@ -282,8 +284,8 @@ def __init__(self, chrom, start, end, *args):
     def __str__(self) -> str:
         return f'Chromatin Info Inaccessible in region {self.chrom}:{self.start}-{self.end}'
 
-def extract_info(chrom, start, end, label, genome_pyfasta, bigwigs, target_bam, strand="+", transforms:dict=None):
-    seq = genome_pyfasta[chrom][int(start):int(end)]
+def extract_info(chrom, start, end, label, genome_pyfaidx, bigwigs, target_bam, strand="+", transforms:dict=None):
+    seq = genome_pyfaidx[chrom][int(start):int(end)].seq
     if strand=="-":
         seq = rev_comp(seq)
     seq_array = dna2OneHot(seq)

seqchromloader/writer.py

@@ -12,7 +12,7 @@
 from collections import defaultdict
 from multiprocessing import Pool
 
-import pyfasta
+import pyfaidx
 import pysam
 import pyBigWig
 import webdataset as wds
@@ -53,7 +53,7 @@ def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
                         compress=True, 
                         numProcessors=1,
                         transforms=None,
-                        braceexpand=True,
+                        braceexpand=False,
                         DALI=False):
     """
     Given coordinates dataframe, extract the sequence and chromatin signal, save in webdataset format
@@ -85,7 +85,7 @@ def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
     # split coordinates and assign chunks to workers
     num_chunks = math.ceil(len(coords) / 7000)
     chunks = np.array_split(coords, num_chunks)
-
+    
     # freeze the common parameters
     ## create a scaler to get statistics for normalizing chromatin marks input
     ## also create a multiprocessing lock
@@ -99,18 +99,19 @@ def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
                                                     DALI=DALI)
 
     count_of_digits = 0
-    while num_chunks > 0:
-       num_chunks = int(num_chunks/10)
+    nc = num_chunks
+    while nc > 0:
+       nc = int(nc/10)
        count_of_digits += 1
 
     pool = Pool(numProcessors)
     res = pool.starmap_async(dump_data_worker_freeze, zip(chunks, [outprefix + "_" + format(i, f'0{count_of_digits}d') for i in range(num_chunks)]))
     files = res.get()
 
     if braceexpand:
-        begin = f'0{count_of_digits}d'.format(0)
-        end = f'0{count_of_digits}d'.format(range(num_chunks)[-1])
-        return f"outprefix_{{{begin}...{end}}}.tar.gz" if compress else f"outprefix_{{{begin}...{end}}}.tar"
+        begin = format(0, f'0{count_of_digits}d')
+        end = format(range(num_chunks)[-1], f'0{count_of_digits}d')
+        return os.path.join(outdir, f"{outprefix}_{{{begin}..{end}}}.tar.gz" if compress else f"{outprefix}_{{{begin}...{end}}}.tar")
     else:
         return files
 
@@ -124,7 +125,7 @@ def dump_data_webdataset_worker(coords,
                                 transforms=None,
                                 DALI=False):
     # get handlers
-    genome_pyfasta = pyfasta.Fasta(fasta)
+    genome_pyfaidx = pyfaidx.Fasta(fasta)
     bigwigs = [pyBigWig.open(bw) for bw in bigwig_files]
     target_pysam = pysam.AlignmentFile(target_bam) if target_bam is not None else None
 
@@ -141,7 +142,7 @@ def dump_data_webdataset_worker(coords,
                 item.start,
                 item.end,
                 item.label,
-                genome_pyfasta=genome_pyfasta,
+                genome_pyfaidx=genome_pyfaidx,
                 bigwigs=bigwigs,
                 target_bam=target_pysam,
                 strand=item.strand,

setup.py

@@ -20,7 +20,7 @@
     # eg: 1.0.0, 1.0.1, 3.0.2, 5.0-beta, etc.
     # You CANNOT upload two versions of your package with the same version number
     # This field is REQUIRED
-    version="0.5.1",
+    version="0.5.2",
 
     # The packages that constitute your project.
     # For my project, I have only one - "pydash".
@@ -37,7 +37,7 @@
         'numpy',
         'pandas',
         'webdataset>=0.2.0',
-        'pyfasta>=0.5.0',
+        'pyfaidx>=0.7.0',
         'pybedtools>=0.9.0',
         'pysam>=0.19.0',
         'pybigwig>=0.3.0',

tests/test_writer_loader.py

@@ -1,6 +1,5 @@
 import os
 import sys
-sys.path.insert(0, "./")
 import numpy as np
 import pandas as pd
 from seqchromloader import SeqChromDatasetByDataFrame, SeqChromDatasetByBed, SeqChromDatasetByWds, SeqChromDataModule
@@ -68,7 +67,7 @@ def test_random_coords(self):
         coords_incl = random_coords(genome="mm10", incl=interval)
         coords_excl = random_coords(genome="mm10", excl=interval)
 
-        self.assertTrue(BedTool().from_dataframe(coords_incl).intersect(interval).count()==coords_incl.size)
+        self.assertTrue(BedTool().from_dataframe(coords_incl).intersect(interval).count()==len(coords_incl))
         self.assertTrue(BedTool().from_dataframe(coords_excl).intersect(interval).count()==0)
 
     def test_chop_genome(self):
@@ -77,9 +76,10 @@ def test_chop_genome(self):
                                                       'end': [50000, 20000]}))
         coords_incl = chop_genome(chroms=["chr2", "chr12"], genome="mm10", stride=1000, l=500, incl=interval)
         coords_excl = chop_genome(chroms=["chr2", "chr12"], genome="mm10", stride=1000, l=500, excl=interval)
-        self.assertTrue(np.all([coords_incl.start.iloc[i] - coords_incl.start.iloc[i-1] for i in range(1, len(coords_incl))]==1000))
-        self.assertTrue(np.all([coords_excl.start.iloc[i] - coords_excl.start.iloc[i-1] for i in range(1, len(coords_excl))]==1000))
-        self.assertTrue(BedTool().from_dataframe(coords_incl).intersect(interval).count()==coords_incl.size)
+        for c in ['chr2', 'chr12']:
+            df = coords_incl[coords_incl.chrom==c]
+            self.assertTrue(np.all([df.start.iloc[i] - df.start.iloc[i-1] == 1000 for i in range(1, len(df))]))
+        self.assertTrue(BedTool().from_dataframe(coords_incl).intersect(interval).count()==len(coords_incl))
         self.assertTrue(BedTool().from_dataframe(coords_excl).intersect(interval).count()==0)
 
     def test_writer(self):
@@ -99,8 +99,18 @@ def test_writer(self):
                     outdir=self.tempdir,
                     outprefix='test',
                     compress=True,
-                    numProcessors=5)
+                    numProcessors=2)
         self.assertIsFile(os.path.join(self.tempdir, "test_0.tar.gz"))
+        wds_files = dump_data_webdataset(huge_coords, 
+                    genome_fasta='data/sample.fa', 
+                    bigwig_filelist=['data/sample.bw'],
+                    target_bam='data/sample.bam',
+                    outdir=self.tempdir,
+                    outprefix='test',
+                    compress=True,
+                    numProcessors=2,
+                    braceexpand=True)
+        self.assertTrue(wds_files == os.path.join(self.tempdir, "test_{0..1}.tar.gz"))
 
         ds = wds.DataPipeline(
             wds.SimpleShardList([os.path.join(self.tempdir, "test_0.tar.gz")]),
@@ -251,4 +261,4 @@ def test_target_transform(target):
     return target * 3
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main(verbosity=2)