add pybigtools backend

Author: yztxwd

seqchromloader/loader.py

@@ -10,7 +10,7 @@
 import random
 import pysam
 import pyfaidx
-import pyBigWig
+import pybigtools
 import numpy as np
 import pandas as pd
 import webdataset as wds
@@ -19,7 +19,7 @@
 from torch.utils.data import Dataset, IterableDataset, DataLoader
 from pybedtools import BedTool
 
-from seqchromloader import utils
+from seqchromloader import utils, config
 
 logger = logging.getLogger(__name__)
 
@@ -139,7 +139,7 @@ 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_pyfaidx = pyfaidx.Fasta(self.genome_fasta)
-        self.bigwigs = [pyBigWig.open(bw) for bw in self.bigwig_filelist] if self.bigwig_filelist is not None else None
+        self.bigwigs = [utils.BigWig(bw_path) for bw_path in self.bigwig_filelist] if self.bigwig_filelist is not None else None
         if self.target_bam is not None:
             if isinstance(self.target_bam, list):
                 self.target_pysam = [pysam.AlignmentFile(b) for b in self.target_bam]

seqchromloader/utils.py

@@ -8,14 +8,49 @@
 import logging
 import pysam
 import pyBigWig
+import pybigtools
 from Bio import motifs
 from pyfaidx import Fasta
 from multiprocessing import Pool
 from pybedtools import Interval, BedTool
 from pybedtools.helpers import chromsizes
 
+from seqchromloader import config
+
 logger = logging.getLogger(__name__)
 
+class BigWig():
+    def __init__(self, bw_path, backend='pyBigWig'):
+        if backend == 'pyBigWig':
+            self.bw = pyBigWig.open(bw_path)
+        else:
+            self.bw = pybigtools.open(bw_path)
+        self.backend = backend
+
+    def intervals(self, chrom):
+        if self.backend == 'pyBigWig':
+            return self.bw.intervals(chrom)
+        else:
+            return self.bw.records(chrom)
+
+    def stats(self, chrom, type='mean'):
+        if self.backend == 'pyBigWig':
+            return self.bw.stats(chrom, type=type, exact=True)[0]
+        else:
+            return self.bw.values(chrom, missing=np.nan, bins=1, exact=True, summary='mean')[0].item()
+
+    def values(self, chrom, start, end, missing=0):
+        if self.backend == 'pyBigWig':
+            return np.nan_to_num(self.bw.values(chrom, start, end)).astype(np.float32)
+        else:
+            return self.bw.values(chrom, start, end, missing=0.).astype(np.float32)
+    
+    def chroms(self):
+        return self.chroms()
+
+    def close(self):
+        self.bw.close()
+
 def get_genome_sizes(gs=None, genome=None, to_filter=None, to_keep=None):
     """
     Loads the genome sizes file, filter or keep chromosomes
@@ -372,7 +407,7 @@ def compute_mean_std_bigwig(bigwig):
     :type bigwig: str
     :rtype: (mean, stddev)
     """
-    bw = pyBigWig.open(bigwig)
+    bw = BigWig(bigwig)
 
     # get chrom length list
     chroms = bw.chroms()
@@ -485,7 +520,7 @@ def extract_bw(chrom, start, end, strand, bigwigs):
     chroms_array = []
     try:
         for idx, bigwig in enumerate(bigwigs):
-            c = (np.nan_to_num(bigwig.values(chrom, start, end))).astype(np.float32)
+            c = bigwig.values(chrom, start, end)
             if strand=="-":
                 c = c[::-1] 
             chroms_array.append(c)
@@ -508,9 +543,9 @@ def extract_dnaOneHot(chrom, start, end, strand, genome_pyfaidx):
 def extract_single_target(chrom, start, end, strand, target):
     if isinstance(target, pysam.AlignmentFile):
         target_array = np.array(target.count(chrom, start, end), dtype=np.float32)[np.newaxis]
-    elif isinstance(target, pyBigWig.pyBigWig):
+    elif isinstance(target, BigWig):
         try:
-            target_array = np.nan_to_num(target.values(chrom, start, end)).astype(np.float32)
+            target_array = target.values(chrom, start, end)
             if strand=="-":
                 target_array = target_array[::-1]
         except RuntimeError as e:

seqchromloader/writer.py

@@ -14,11 +14,10 @@
 
 import pyfaidx
 import pysam
-import pyBigWig
 import webdataset as wds
 
-from . import utils
-from .loader import _SeqChromDatasetByWds
+from seqchromloader import utils
+from seqchromloader.loader import _SeqChromDatasetByWds
 
 logger = logging.getLogger(__name__)
 
@@ -148,17 +147,17 @@ def dump_data_webdataset_worker(coords,
                                 ):
     #get handlers
     genome_pyfaidx = pyfaidx.Fasta(fasta)
-    bigwigs = [pyBigWig.open(bw) for bw in bigwig_files] if bigwig_files is not None else None
+    bigwigs = [utils.BigWig(bw) for bw in bigwig_files] if bigwig_files is not None else None
     if target_bam is not None:
         if isinstance(target_bam, list):
             target = [pysam.AlignmentFile(b) for b in target_bam]
         else:
             target = pysam.AlignmentFile(target_bam)
     elif target_bw is not None:
         if isinstance(target_bw, list):
-            target = [pyBigWig.open(b) for b in target_bw]
+            target = [utils.BigWig(b) for b in target_bw]
         else:
-            target = pyBigWig.open(target_bw)
+            target = utils.BigWig(target_bw)
     else:
         target = None
 

tests/test_writer_loader.py

@@ -18,6 +18,10 @@
 from pyfaidx import Fasta
 from pyjaspar import jaspardb
 
+import seqchromloader
+
+seqchromloader.config.set_bigwig_backend('pybigtools')
+
 class Test(unittest.TestCase):
     def setUp(self) -> None:
         pass