seqchromloader 0.3.0__tar.gz → 0.4.0__tar.gz

{seqchromloader-0.3.0 → seqchromloader-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: seqchromloader
-Version: 0.3.0
+Version: 0.4.0
 Summary: Sequence and chromatin dataloader for deep learning
 Home-page: https://github.com/yztxwd/seqchromloader
 Author-email: yztxwd@gmail.com

{seqchromloader-0.3.0 → seqchromloader-0.4.0}/seqchromloader/__init__.py

@@ -1,2 +1,2 @@
 from .loader import SeqChromDatasetByDataFrame, SeqChromDatasetByBed, SeqChromDatasetByWds, SeqChromDataModule
-from .writer import dump_data_webdataset
+from .writer import dump_data_webdataset, convert_data_webdataset

{seqchromloader-0.3.0 → seqchromloader-0.4.0}/seqchromloader/loader.py

@@ -57,12 +57,13 @@ class _SeqChromDatasetByWds(IterableDataset):
     :param transforms: A dictionary of functions to transform the output data, accepted keys are **["seq", "chrom", "target", "label"]**
     :type transforms: dict of functions
     """
-    def __init__(self, wds, transforms:dict=None, rank=0, world_size=1):
+    def __init__(self, wds, transforms:dict=None, rank=0, world_size=1, keep_key=False):
         self.wds = wds
         self.transforms = transforms
 
         self.rank = rank
         self.world_size = world_size
+        self.keep_key = keep_key
 
     def initialize(self):
         # this function will be called by worker_init_function in DataLoader
@@ -85,7 +86,10 @@ class _SeqChromDatasetByWds(IterableDataset):
         if self.transforms is not None: 
             pipeline.append(wds.map_dict(**self.transforms))
 
-        pipeline.append(wds.to_tuple("seq", "chrom", "target", "label"))
+        if self.keep_key:
+            pipeline.append(wds.to_tuple("__key__", "seq", "chrom", "target", "label"))
+        else:
+            pipeline.append(wds.to_tuple("seq", "chrom", "target", "label"))
             
         ds = wds.DataPipeline(*pipeline)
 

{seqchromloader-0.3.0 → seqchromloader-0.4.0}/seqchromloader/writer.py

@@ -17,14 +17,43 @@ import pysam
 import pyBigWig
 import webdataset as wds
 
-from seqchromloader import utils
+from . import utils
+from .loader import _SeqChromDatasetByWds
 
+def convert_data_webdataset(wds_in, wds_out, transforms=None, compress=False):
+    """
+    Transform the provided webdataset
+    
+    :param wds_in: input webdataset file
+    :type wds_in: string
+    :param wds_out: output webdataset file
+    :type wds_out: string
+    :param transforms: A dictionary of functions to transform the output data, accepted keys are *["seq", "chrom", "target", "label"]*
+    :type transforms: dict of functions
+    :param compress: whether to compress the output file
+    :type compress: boolean
+    """
+    
+    ds = _SeqChromDatasetByWds(wds_in, transforms=transforms, keep_key=True)
+    sink = wds.TarWriter(wds_out, compress=compress)
+    for (key, seq, chrom, target, label) in ds:
+        feature_dict = defaultdict()
+        feature_dict["__key__"] = key
+        
+        feature_dict["seq.npy"] = seq
+        feature_dict["chrom.npy"] = chrom
+        feature_dict["target.npy"] = target
+        feature_dict["label.npy"] = label
+        sink.write(feature_dict)
+    sink.close()
+    
 def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
                         target_bam=None, 
                         outdir="dataset/", outprefix="seqchrom", 
                         compress=True, 
                         numProcessors=1,
-                        transforms=None):
+                        transforms=None,
+                        DALI=False):
     """
     Given coordinates dataframe, extract the sequence and chromatin signal, save in webdataset format
 
@@ -46,6 +75,8 @@ def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
     :type compress: boolean
     :param numProcessors: number of processors
     :type numProcessors: int
+    :param DALI: Set to True if you want to use the dataset for NVIDIA DALI, it would save all arrays in bytes, which results in losing the array shape info
+    :param DALI: boolean
     """
 
     # split coordinates and assign chunks to workers
@@ -61,10 +92,16 @@ def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
                                                     target_bam=target_bam,
                                                     compress=compress,
                                                     outdir=outdir,
-                                                    transforms=transforms)
+                                                    transforms=transforms,
+                                                    DALI=DALI)
+    
+    count_of_digits = 0
+    while num_chunks > 0:
+       num_chunks = int(num_chunks/10)
+       count_of_digits += 1
 
     pool = Pool(numProcessors)
-    res = pool.starmap_async(dump_data_worker_freeze, zip(chunks, [outprefix + "_" + str(i) for i in range(num_chunks)]))
+    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()
 
     return files
@@ -76,7 +113,8 @@ def dump_data_webdataset_worker(coords,
                                 target_bam=None, 
                                 outdir="dataset/", 
                                 compress=True,
-                                transforms=None):
+                                transforms=None,
+                                DALI=False):
     # get handlers
     genome_pyfasta = pyfasta.Fasta(fasta)
     bigwigs = [pyBigWig.open(bw) for bw in bigwig_files]
@@ -103,11 +141,17 @@ def dump_data_webdataset_worker(coords,
             )
         except utils.BigWigInaccessible as e:
             continue
-
-        feature_dict["seq.npy"] = feature['seq']
-        feature_dict["chrom.npy"] = feature['chrom']
-        feature_dict["target.npy"] = feature['target']
-        feature_dict["label.npy"] = feature['label']
+        
+        if not DALI:
+            feature_dict["seq.npy"] = feature['seq']
+            feature_dict["chrom.npy"] = feature['chrom']
+            feature_dict["target.npy"] = feature['target']
+            feature_dict["label.npy"] = feature['label']
+        else:
+            feature_dict["seq.npy"] = feature['seq'].tobytes()
+            feature_dict["chrom.npy"] = feature['chrom'].tobytes()
+            feature_dict["target.npy"] = feature['target'].tobytes()
+            feature_dict["label.npy"] = feature['label'].tobytes()
 
         sink.write(feature_dict)
 

{seqchromloader-0.3.0 → seqchromloader-0.4.0}/seqchromloader.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: seqchromloader
-Version: 0.3.0
+Version: 0.4.0
 Summary: Sequence and chromatin dataloader for deep learning
 Home-page: https://github.com/yztxwd/seqchromloader
 Author-email: yztxwd@gmail.com

{seqchromloader-0.3.0 → seqchromloader-0.4.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     # 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.3.0",
+    version="0.4.0",
 
     # The packages that constitute your project.
     # For my project, I have only one - "pydash".