use new reader to load data in tracks

Author: Nanguage

coolbox/core/track/arcs/base.py

@@ -1,6 +1,9 @@
+import typing as T
+
 import pandas as pd
 
 from coolbox.utilities import GenomeRange
+from coolbox.utilities.reader.tab import get_indexed_tab_reader
 from coolbox.core.track.base import Track
 from .plot import PlotContacts
 
@@ -83,18 +86,24 @@ def __init__(self, **kwargs):
         properties = ArcsBase.DEFAULT_PROPERTIES.copy()
         properties.update(kwargs)
         super().__init__(properties)
+        self.reader = get_indexed_tab_reader(self.file)
 
-    def fetch_plot_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
+    def fetch_data(
+            self,
+            gr: GenomeRange,
+            gr2: T.Optional[GenomeRange] = None,
+            **kwargs) -> pd.DataFrame:
         """
-
         Returns
         -------
         intervals : pandas.core.frame.DataFrame
             Can be two types:
             1: with columns: ['pos1', 'pos2', 'score'] 'score' is optional
             2: with columns: ['start1', 'end1', 'start2', 'end2', 'score'] 'score' is optional
         """
-        return self.fetch_data(gr, **kwargs)
+        open_region = self.properties.get("open_region") in ["yes", True]
+        df = self.reader.query_var_chr(gr, second=gr2, open_region=open_region, **kwargs)
+        return df
 
     def plot(self, ax, gr: GenomeRange, **kwargs):
         """

coolbox/core/track/arcs/bedpe.py

@@ -1,12 +1,10 @@
 import pandas as pd
 
 from .base import ArcsBase
-from .fetch import FetchParix
-from coolbox.utilities.bed import process_bedpe
 from coolbox.utilities.genome import GenomeRange
 
 
-class BEDPE(ArcsBase, FetchParix):
+class BEDPE(ArcsBase):
     """
     Arcs track from .bedpe file.
 
@@ -32,26 +30,6 @@ def __init__(self, file, **kwargs):
             **kwargs
         })
         super().__init__(**properties)
-        self.bgz_file = process_bedpe(file)
-
-    def fetch_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
-        # filter peaks manually for hicpeaks style in fetch_plot_data
-        df = self.fetch_intervals(self.bgz_file, gr, kwargs.get('gr2'))
-        # TODO the returned df has no named columns, may cause error
-        if len(df) == 0:
-            return df
-
-        columns = list(df.columns)
-        for i, col in enumerate(self.FIELDS):
-            if i >= len(columns):
-                break
-            columns[i] = col
-        df.columns = columns
-        for col in ['start1', 'end1', 'start2', 'end2']:
-            df[col] = df[col].astype(int)
-        if 'score' in df.columns:
-            df['score'] = df['score'].astype(float)
-        return df
 
     def fetch_plot_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
         df = self.fetch_data(gr, **kwargs)

coolbox/core/track/arcs/fetch.py

@@ -1,12 +1,8 @@
-import pandas as pd
-
 from .base import ArcsBase
-from .fetch import FetchParix
-from coolbox.utilities.bed import process_pairs
 from coolbox.utilities.genome import GenomeRange
 
 
-class Pairs(ArcsBase, FetchParix):
+class Pairs(ArcsBase):
     """
     Arcs track from .pairs file.
 
@@ -27,22 +23,4 @@ def __init__(self, file, **kwargs):
             'file': file,
             **kwargs
         })
-        super().__init__(**properties)
-        self.bgz_file = process_pairs(file)
-
-    def fetch_data(self, gr: GenomeRange, **kwargs):
-        # filter peaks manually in peaks style
-        df = self.fetch_intervals(self.bgz_file, gr, kwargs.get('gr2'))
-        # TODO the returned df has no named columns, may cause error
-        if len(df) == 0:
-            return df
-
-        columns = list(df.columns)
-        for i, col in enumerate(self.FIELDS):
-            if i >= len(columns):
-                break
-            columns[i] = col
-        df.columns = columns
-        for col in ['pos1', 'pos2']:
-            df[col] = df[col].astype(int)
-        return df
+        super().__init__(**properties)

coolbox/core/track/arcs/pairs.py

@@ -5,7 +5,7 @@
 from coolbox.utilities import (
     get_logger, GenomeRange, split_genome_range
 )
-from coolbox.utilities.bam import process_bam, query_bam
+from coolbox.utilities.reader.tab import get_indexed_tab_reader
 from .base import Track
 
 log = get_logger(__name__)
@@ -39,7 +39,7 @@ def __init__(self, file, **kwargs):
         })
         properties.update(kwargs)
         super().__init__(properties)
-        self.indexed_bam = process_bam(file)
+        self.reader = get_indexed_tab_reader(file)
 
     def fetch_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
         """
@@ -53,7 +53,7 @@ def fetch_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
             columns = ["qname", "flag", "rname", "pos", "mapq", "cigar",
                       "rnext", "pnext", "tlen", "seq", "qual", "options"]
         """
-        return self.fetch_intervals(gr)
+        return self.reader.query_var_chr(gr)
 
     def plot(self, ax, gr: GenomeRange, **kwargs):
         self.plot_align(ax, gr)
@@ -85,22 +85,3 @@ def plot_align(self, ax, gr: GenomeRange):
             with_ruler=False,
             draw_line=False
         )
-
-    def fetch_intervals(self, genome_range: GenomeRange):
-        chrom, start, end = split_genome_range(genome_range)
-        rows = [
-            row_items
-            for row_items in query_bam(
-                self.indexed_bam, chrom, start, end, split=True
-            )
-        ]
-
-        # https://samtools.github.io/hts-specs/SAMv1.pdf
-        fields = ["qname", "flag", "rname", "pos", "mapq", "cigar",
-                  "rnext", "pnext", "tlen", "seq", "qual", "options"]
-        df = pd.DataFrame(rows, columns=fields)
-        if df.shape[0] > 0:
-            df['flag'] = df['flag'].astype(int)
-            df['pos'] = df['pos'].astype(int)
-            df['mapq'] = df['mapq'].astype(int)
-        return df

coolbox/core/track/bam.py

@@ -4,7 +4,7 @@
 import matplotlib
 
 from coolbox.utilities import get_logger
-from coolbox.utilities.bed import build_bed_index
+from coolbox.utilities.reader.tab import get_indexed_tab_reader
 from coolbox.utilities.genome import GenomeRange
 from coolbox.core.track.base import Track
 
@@ -52,7 +52,7 @@ def __init__(self, file, **kwargs):
             **kwargs
         })
         super().__init__(properties)
-        self.bgz_file = build_bed_index(file)
+        self.reader = get_indexed_tab_reader(file)
 
     def fetch_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
         """
@@ -135,18 +135,21 @@ def get_rgb_and_edge_color(self, bed):
     @staticmethod
     def infer_bed_type(df: pd.DataFrame) -> Union[str, None]:
         #  bed_type of dataframe are store in dataframe's __dict__ in FetchBed.fetch_intervals
-        if 'bed_type' in df.__dict__:
-            bed_type = df.bed_type
-        else:
-            bed_types = {
-                12: 'bed12',
-                9: 'bed9',
-                6: 'bed6',
-                3: 'bed3'
-            }
-            num_col = len(df.columns)
-            bed_type = bed_types[num_col] if num_col in bed_types else 'bed3'
-            if bed_type == 'bed3' and num_col < 3:
-                raise ValueError(f"Invalid dataframe for bed3 with columns: {df.columns}")
+        bed_types = {
+            12: 'bed12',
+            9: 'bed9',
+            6: 'bed6',
+            3: 'bed3'
+        }
+        num_col = len(df.columns)
+        bed_type = bed_types[num_col] if num_col in bed_types else 'bed3'
+        if bed_type == 'bed3' and num_col < 3:
+            raise ValueError(f"Invalid dataframe for bed3 with columns: {df.columns}")
         return bed_type
 
+    def fetch_intervals(self, gr: GenomeRange) -> pd.DataFrame:
+        """
+        Fetch intervals within input chromosome range.
+        """
+        df = self.reader.query_var_chr(gr)
+        return df

coolbox/core/track/bed/base.py

@@ -1,15 +1,17 @@
-from coolbox.core.track.bed.fetch import FetchBed
 from coolbox.utilities import (
     get_logger
 )
 from coolbox.utilities.genome import GenomeRange
 from .base import BedBase
 from .plot import PlotGenes
 
+import pandas as pd
+
+
 log = get_logger(__name__)
 
 
-class BED(BedBase, PlotGenes, FetchBed):
+class BED(BedBase, PlotGenes):
     """
     Bed Track for plotting 1d intervals data from .bed file.
     The input bed file can be bed3/bed6/bed9/bed12

coolbox/core/track/bed/bed.py

@@ -1,5 +1,5 @@
 from .base import HistBase, HistData, GenomeRange
-from coolbox.utilities.bam import process_bam, coverage_by_samtools
+from coolbox.utilities.reader.tab import process_bam, coverage_by_samtools
 
 
 class BAMCov(HistBase):

coolbox/core/track/bed/fetch.py

@@ -1,10 +1,7 @@
 import pandas as pd
 
-from coolbox.utilities import (
-    change_chrom_names,
-    GenomeRange, get_logger,
-)
-from coolbox.utilities.bed import tabix_query, build_bedgraph_bgz
+from coolbox.utilities import GenomeRange, get_logger
+from coolbox.utilities.reader.tab import get_indexed_tab_reader
 from .base import HistBase
 
 log = get_logger(__name__)
@@ -32,7 +29,7 @@ def __init__(self, file, **kwargs):
             **kwargs
         })
         super().__init__(**properties)
-        self.bgz_file = build_bedgraph_bgz(file)
+        self.reader = get_indexed_tab_reader(file)
 
     def fetch_plot_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
         itv_df = self.fetch_data(gr, **kwargs)
@@ -41,18 +38,4 @@ def fetch_plot_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
         return itv_df
 
     def fetch_data(self, gr: GenomeRange, **kwargs) -> pd.DataFrame:
-        rows = self.load(gr)
-        if len(rows) == 0:
-            gr.chrom = change_chrom_names(gr.chrom)
-            rows = self.load(gr)
-
-        return pd.DataFrame(rows, columns=['chromsome', 'start', 'end', 'score'])
-
-    def load(self, genome_range):
-        gr = genome_range
-        return [
-            [it[0], int(it[1]), int(it[2]), float(it[3])]
-            for it in tabix_query(
-                self.bgz_file, gr.chrom, gr.start, gr.end, split=True
-            )
-        ]
+        return self.reader.query_var_chr(gr)