Merge pull request #48 from jksr/master

ballc format support

Author: lhqing

ALLCools/_allc_to_region_count.py

@@ -135,6 +135,7 @@ def _map_to_sparse_chrom_bin(site_bed, out_bed, chrom_size_path, bin_size=500):
 )
 def allc_to_region_count(
     allc_path: str,
+    cmeta_path: str,
     output_prefix: str,
     chrom_size_path: str,
     mc_contexts: List[str],
@@ -215,6 +216,7 @@ def allc_to_region_count(
     strandness = "split" if split_strand else "both"
     output_paths_dict = extract_allc(
         allc_path=allc_path,
+        cmeta_path=cmeta_path,
         output_prefix=output_prefix,
         mc_contexts=mc_contexts,
         strandness=strandness,
@@ -332,10 +334,11 @@ def batch_allc_to_region_count(
 
     with ProcessPoolExecutor(cpu) as executor:
         futures = {}
-        for cell_id, allc_path in allc_series.iteritems():
+        for cell_id, (allc_path, cmeta_path) in allc_series.iterrows():
             future = executor.submit(
                 allc_to_region_count,
                 allc_path=allc_path,
+                cmeta_path=cmeta_path,
                 output_prefix=str(output_dir / cell_id),
                 chrom_size_path=chrom_size_path,
                 mc_contexts=mc_contexts,

ALLCools/_extract_allc.py

@@ -219,6 +219,7 @@ def _extract_allc_parallel(
 )
 def extract_allc(
     allc_path: str,
+    cmeta_path: str,
     output_prefix: str,
     mc_contexts: Union[str, list],
     chrom_size_path: str,
@@ -346,7 +347,7 @@ def extract_allc(
 
     # split file first
     # strandness function
-    with open_allc(allc_path, region=region) as allc:
+    with open_allc(allc_path, region=region, cmeta_path=cmeta_path) as allc:
         if strandness == "Split":
             for line in allc:
                 cur_line = line.split("\t")

ALLCools/_open.py

@@ -206,7 +206,80 @@ def _raise_if_error(self):
         """
         Raise an exception if the gzip process has exited with an error.
 
-        Raise IOError if process is not running anymore and the
+        Raise OSError if process is not running anymore and the
+        exit code is nonzero.
+        """
+        return_code = self.process.poll()
+        if return_code is not None and return_code != 0:
+            message = self._stderr.read().strip()
+            raise OSError(message)
+
+    def read(self, *args):
+        data = self._file.read(*args)
+        if len(args) == 0 or args[0] <= 0:
+            # wait for process to terminate until we check the exit code
+            self.process.wait()
+        self._raise_if_error()
+        return data
+
+
+class PipedBAllCReader(Closing):
+    def __init__(self, path, cmeta_path=None, region=None, mode="r"):
+        if mode not in ("r"):
+            raise ValueError(f"Mode can only be 'r' for ballc file")
+        if cmeta_path is None:
+            raise NotImplementedError
+        cmeta_path = pathlib.Path(cmeta_path).resolve()
+        if not cmeta_path.exists():
+            raise FileNotFoundError(f"{cmeta_path} does not exist.")
+        cmeta_path = str(cmeta_path)
+
+        if region is None:
+            self.process = Popen(
+                ["ballcools", "query", path, '"*"', "-c", cmeta_path],
+                # stdout=PIPE, stderr=PIPE, encoding="utf8")
+                stdout=PIPE,
+                stderr=PIPE,
+                encoding=None,
+            )
+        else:
+            self.process = Popen(
+                ["ballcools", "query", path] + region.split(" ") + ["-c", cmeta_path],
+                stdout=PIPE,
+                stderr=PIPE,
+                # encoding="utf8",)
+                encoding=None,
+            )
+
+        self.name = path
+        self._file = self.process.stdout
+        self._stderr = self.process.stderr
+        self.closed = False
+        # Give ballcools a little bit of time to report any errors
+        # (such as a non-existing file)
+        time.sleep(0.01)
+        self._raise_if_error()
+
+    def close(self):
+        self.closed = True
+        return_code = self.process.poll()
+        if return_code is None:
+            # still running
+            self.process.terminate()
+        self._raise_if_error()
+
+    def __iter__(self):
+        for line in self._file:
+            yield line.decode("utf-8")
+        self.process.wait()
+        self._raise_if_error()
+
+    def readline(self):
+        return self._file.readline().decode("utf-8")
+
+    def _raise_if_error(self):
+        """
+        Raise OSError if process is not running anymore and the
         exit code is nonzero.
         """
         return_code = self.process.poll()
@@ -275,7 +348,7 @@ def readline(self):
         return self.file.readline()
 
     def _raise_if_error(self):
-        """Raise IOError if process is not running anymore and the exit code is nonzero."""
+        """Raise OSError if process is not running anymore and the exit code is nonzero."""
         return_code = self.process.poll()
         if return_code is not None and return_code != 0:
             message = self._stderr.read().strip()
@@ -382,7 +455,11 @@ def open_gz(file_path, mode="r", compresslevel=3, threads=1, region=None):
             return gzip.open(file_path, mode, compresslevel=compresslevel)
 
 
-def open_allc(file_path, mode="r", compresslevel=3, threads=1, region=None):
+def open_ballc(file_path, mode="r", compresslevel=None, threads=None, region=None, cmeta_path=None):
+    return PipedBAllCReader(file_path, cmeta_path=cmeta_path, region=region, mode="r")
+
+
+def open_allc(file_path, mode="r", compresslevel=3, threads=1, region=None, cmeta_path=None):
     """
     Open a .allc file.
 
@@ -422,7 +499,7 @@ def open_allc(file_path, mode="r", compresslevel=3, threads=1, region=None):
         raise ValueError(f"mode '{mode}' not supported")
     if compresslevel not in range(1, 10):
         raise ValueError("compresslevel must be between 1 and 9")
-    if region is not None:
+    if (region is not None) and (not file_path.endswith(".ballc")):
         # unzipped file
         if not file_path.endswith("gz"):
             raise ValueError(f"File must be compressed by bgzip to use region query. File path {file_path}")
@@ -435,8 +512,10 @@ def open_allc(file_path, mode="r", compresslevel=3, threads=1, region=None):
         if not os.path.exists(file_path + ".tbi"):
             raise FileNotFoundError("region query provided, but .tbi index not found")
 
-    if file_path.endswith("gz"):
+    if file_path.endswith(".gz"):
         return open_gz(file_path, mode, compresslevel, threads, region=region)
+    elif file_path.endswith(".ballc"):
+        return open_ballc(file_path, region=region, cmeta_path=cmeta_path)
     else:
         return open(file_path, mode)
 

ALLCools/count_matrix/mcds.py

@@ -331,11 +331,17 @@ def generate_mcds(
     dtype = parse_dtype(dtype)
 
     if isinstance(allc_table, str):
-        allc_series = pd.read_csv(allc_table, header=None, index_col=0, squeeze=True, sep="\t")
-        if not isinstance(allc_series, pd.Series):
-            raise ValueError("allc_table malformed, should only have 2 columns, 1. file_uid, 2. file_path")
+        allc_series = pd.read_csv(allc_table, header=None, index_col=0, sep="\t")
+        if allc_series.shape[1] == 1:
+            allc_series[2] = ""
+        elif allc_series.shape[1] == 2:
+            pass
+        else:
+            raise ValueError(
+                "allc_table malformed, should have 2 or 3 columns, 1. file_uid, 2. file_path, 3. cmeta_path if ballc format"
+            )
     else:
-        allc_series = allc_table.dropna()
+        allc_series = allc_table[~allc_table[1].isna()]
     if allc_series.index.duplicated().sum() != 0:
         raise ValueError("allc_table file uid have duplicates (1st column)")
 
@@ -363,7 +369,7 @@ def generate_mcds(
         else:
             if n_matched < rna_series.size:
                 print(f"{rna_series.size - n_matched} file_uid in RNA table do not exist in ALLC table")
-            if n_matched < allc_series.size:
+            if n_matched < allc_series.shape[0]:
                 print(f"{allc_series.size - n_matched} file_uid in ALLC table do not exist in RNA table")
 
         # reindex allc and rna series, this may create NaN, will be dropped
@@ -372,19 +378,19 @@ def generate_mcds(
         rna_series = rna_series.reindex(union_index)
 
     # if allc files exceed max_per_mcds, save them into chunks
-    if allc_series.size > max_per_mcds:
-        mcds_n_chunk = ceil(allc_series.size / max_per_mcds)
-        chunk_size = ceil(allc_series.size / mcds_n_chunk)
+    if allc_series.shape[0] > max_per_mcds:
+        mcds_n_chunk = ceil(allc_series.shape[0] / max_per_mcds)
+        chunk_size = ceil(allc_series.shape[0] / mcds_n_chunk)
         allc_series_chunks = [
-            allc_series[chunk_start : chunk_start + chunk_size]
-            for chunk_start in range(0, allc_series.size, chunk_size)
+            allc_series.iloc[chunk_start : chunk_start + chunk_size]
+            for chunk_start in range(0, allc_series.shape[0], chunk_size)
         ]
         if rna_table is not None:
             rna_series_chunks = [
                 rna_series[chunk_start : chunk_start + chunk_size]
                 for chunk_start in range(0, rna_series.size, chunk_size)
             ]
-        print(f"Number of file_uids {allc_series.size} > max_per_mcds {max_per_mcds}, ")
+        print(f"Number of file_uids {allc_series.shape[0]} > max_per_mcds {max_per_mcds}, ")
         print(f"will generate MCDS in {len(allc_series_chunks)} chunks.")
 
         # mcds chunks execute sequentially