Generalize variant source in VCF matching

Author: Hoeze

kipoiseq/extractors/vcf.py

@@ -1,9 +1,12 @@
+import abc
 import logging
-from typing import Tuple, Iterable, List, Union
+from typing import Tuple, Iterable, List, Union, Iterator
 from itertools import islice
 from collections import defaultdict
 from kipoiseq.dataclasses import Variant, Interval
+from kipoiseq.variant_source import VariantFetcher
 from kipoiseq.extractors.vcf_query import VariantIntervalQueryable
+from kipoiseq.utils import batch_iter
 
 try:
     from cyvcf2 import VCF
@@ -15,19 +18,11 @@
 ]
 
 
-def _batch_iter(variants: Iterable[Variant], batch_size=10000
-                ) -> Iterable[Iterable[Variant]]:
-    batch = list(islice(variants, batch_size))
-    while batch:
-        yield batch
-        batch = list(islice(variants, batch_size))
-
-
-class MultiSampleVCF(VCF):
+class MultiSampleVCF(VariantFetcher, VCF):
 
     def __init__(self, *args, **kwargs):
         from cyvcf2 import VCF
-        super(MultiSampleVCF, self).__init__(*args, **kwargs, strict_gt=True)
+        VCF.__init__(self, *args, **kwargs, strict_gt=True)
         self.sample_mapping = dict(zip(self.samples, range(len(self.samples))))
 
     def fetch_variants(self, interval, sample_id=None):
@@ -79,7 +74,7 @@ def batch_iter(self, batch_size=10000) -> Iterable[Iterable[Variant]]:
             batch_size: size of each batch.
         """
         variants = iter(self)
-        yield from _batch_iter(variants, batch_size=batch_size)
+        yield from batch_iter(variants, batch_size=batch_size)
 
     def query_variants(self, intervals: List[Interval], sample_id=None,
                        progress=False) -> VariantIntervalQueryable:

kipoiseq/extractors/vcf_matching.py

@@ -1,7 +1,8 @@
-from typing import List
+import os
+from typing import List, Union, Iterable, Iterator
 import pandas as pd
 from kipoiseq.dataclasses import Variant, Interval
-from kipoiseq.extractors import MultiSampleVCF
+from kipoiseq.variant_source import VariantFetcher
 
 try:
     from pyranges import PyRanges
@@ -82,14 +83,58 @@ def intervals_to_pyranges(intervals):
     )
 
 
+class PyrangesVariantFetcher(VariantFetcher):
+
+    def __init__(self, variants: List[Variant]):
+        self.variants = variants
+        self._variants_pr = None
+
+    @property
+    def variants_pr(self):
+        # convert to PyRanges on demand
+        if self._variants_pr is None:
+            self._variants_pr = variants_to_pyranges(self.variants)
+        return self._variants_pr
+
+    def fetch_variants(self, interval: Union[Interval, Iterable[Interval]]) -> Iterator[Variant]:
+        if isinstance(interval, Interval):
+            interval = [interval]
+        # convert interval(s) to PyRanges object
+        interval_pr: PyRanges = intervals_to_pyranges(interval)
+        # join with variants
+        pr_join = interval_pr.join(self.variants_pr, suffix='_variant')
+
+        yield from pr_join.df["variant"]
+
+    def __iter__(self) -> Iterator[Variant]:
+        yield from self.variants
+
+
+class VariantFetcherProxy(VariantFetcher):
+
+    def __init__(self, variant_fetcher: VariantFetcher):
+        self.variant_fetcher = variant_fetcher
+
+    def fetch_variants(self, interval: Union[Interval, Iterable[Interval]]) -> Iterator[Variant]:
+        yield from self.variant_fetcher.fetch_variants(interval)
+
+    def batch_iter(self, batch_size=10000) -> Iterator[List[Variant]]:
+        yield from self.variant_fetcher.batch_iter(batch_size)
+
+    def __iter__(self) -> Iterator[Variant]:
+        yield from self.variant_fetcher
+
+
 class BaseVariantMatcher:
     """
     Base variant intervals matcher
     """
 
     def __init__(
             self,
-            vcf_file: str,
+            vcf_file: str = None,
+            variants: List[Variant] = None,
+            variant_fetcher: VariantFetcher = None,
             gtf_path: str = None,
             bed_path: str = None,
             pranges: PyRanges = None,
@@ -101,7 +146,8 @@ def __init__(
         """
 
         Args:
-          vcf_file: path of vcf file
+          vcf_file: (optional) path of vcf file
+          variants: (optional) readily processed variants
           gtf_path: (optional) path of gtf file contains features
           bed_path: (optional) path of bed file
           pranges: (optional) pyranges object
@@ -110,12 +156,37 @@ def __init__(
             pyranges object. This argument is not valid with intervals.
             Currently unused
         """
-        self.vcf = MultiSampleVCF(vcf_file, lazy=vcf_lazy)
+        self.variant_fetcher = self._read_variants(vcf_file, variants, variant_fetcher, vcf_lazy)
         self.interval_attrs = interval_attrs
         self.pr = self._read_intervals(gtf_path, bed_path, pranges,
                                        intervals, interval_attrs, duplicate_attr=True)
         self.variant_batch_size = variant_batch_size
 
+    @staticmethod
+    def _read_variants(
+            vcf_file=None,
+            variants=None,
+            variant_fetcher=None,
+            vcf_lazy: bool = True,
+    ):
+        if vcf_file is not None:
+            from kipoiseq.extractors import MultiSampleVCF
+            vcf = MultiSampleVCF(vcf_file, lazy=vcf_lazy)
+
+            if os.environ.get('PYTEST_RUNNING', '') == 'true':
+                # Ensure that none of the methods actually uses MultiSampleVCF methods by accident
+                vcf = VariantFetcherProxy(vcf)
+
+            return vcf
+        elif variant_fetcher is not None:
+            assert isinstance(variant_fetcher, VariantFetcher), \
+                "Wrong type of variant fetcher: %s" % type(variant_fetcher)
+            return variant_fetcher
+        elif variants is not None:
+            return PyrangesVariantFetcher(variants)
+        else:
+            raise ValueError("No source of variants was specified!")
+
     @staticmethod
     def _read_intervals(gtf_path=None, bed_path=None, pranges=None,
                         intervals=None, interval_attrs=None, duplicate_attr=False):
@@ -172,7 +243,7 @@ def _read_vcf_pyranges(self, batch_size=10000):
         Args:
           batch_size: size of each batch.
         """
-        for batch in self.vcf.batch_iter(batch_size):
+        for batch in self.variant_fetcher.batch_iter(batch_size):
             yield variants_to_pyranges(batch)
 
     def iter_pyranges(self) -> PyRanges:
@@ -210,30 +281,14 @@ def __iter__(self):
 
 class MultiVariantsMatcher(BaseVariantMatcher):
 
-    def __init__(
-            self,
-            vcf_file,
-            gtf_path=None,
-            pranges=None,
-            intervals=None,
-            interval_attrs=None,
-            vcf_lazy=True,
-            variant_batch_size=10000
-    ):
-        super().__init__(
-            vcf_file,
-            gtf_path=gtf_path,
-            pranges=pranges,
-            intervals=intervals,
-            interval_attrs=interval_attrs,
-            vcf_lazy=vcf_lazy,
-            variant_batch_size=variant_batch_size
-        )
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
         if hasattr(self.pr, 'intervals'):
             self.intervals = self.pr.intervals
         else:
             self.intervals = pyranges_to_intervals(self.pr)
 
     def __iter__(self):
         for i in self.intervals:
-            yield i, self.vcf.fetch_variants(i)
+            yield i, self.variant_fetcher.fetch_variants(i)

kipoiseq/utils.py

@@ -1,7 +1,12 @@
+from itertools import islice
+from typing import Iterable, TypeVar
+
 import numpy as np
 from six import string_types
 
 # alphabets:
+from kipoiseq import Variant
+
 DNA = ["A", "C", "G", "T"]
 RNA = ["A", "C", "G", "U"]
 AMINO_ACIDS = ["A", "R", "N", "D", "B", "C", "E", "Q", "Z", "G", "H",
@@ -36,3 +41,18 @@ def parse_dtype(dtype):
             raise ValueError("Unable to parse dtype: {}. \nException: {}".format(dtype, e))
     else:
         return dtype
+
+
+T = TypeVar('T')
+
+
+def batch_iter(items: Iterable[T], batch_size: int) -> Iterable[Iterable[T]]:
+    # ensure this is an iterator
+    item_iter = iter(items)
+    while True:
+        # create next `batch_size` number of items;
+        batch = list(islice(item_iter, batch_size))
+        if len(batch) == 0:
+            break
+
+        yield batch

kipoiseq/variant_source.py

@@ -0,0 +1,39 @@
+import abc
+from typing import Iterable, Union, Iterator
+
+from kipoiseq import Interval, Variant
+from kipoiseq.utils import batch_iter
+
+
+class VariantFetcher(Iterable, metaclass=abc.ABCMeta):
+    """
+    Base class for all variant-returning data sources.
+    """
+
+    @abc.abstractmethod
+    def fetch_variants(self, interval: Union[Interval, Iterable[Interval]]) -> Iterator[Variant]:
+        """
+        Fetch variants that intersect the provided interval(s).
+
+        :param interval:  One or multiple Interval objects
+        :return: Iterator of Variant objects that intersect `interval`
+        """
+        raise NotImplementedError
+
+    def batch_iter(self, batch_size=10000) -> Iterator[Iterable[Variant]]:
+        """
+        Fetch variants in batches.
+
+        :param batch_size: Number of variants per batch
+        :return: Iterator that yields batches of Variant objects
+        """
+        yield from batch_iter(self, batch_size)
+
+    @abc.abstractmethod
+    def __iter__(self) -> Iterator[Variant]:
+        """
+        Fetch variants in batches.
+
+        :return: Iterator of Variant objects
+        """
+        raise NotImplementedError

tests/extractors/test_vcf_matching.py

@@ -4,7 +4,7 @@
 from kipoiseq.dataclasses import Interval, Variant
 from kipoiseq.extractors.vcf import MultiSampleVCF
 from kipoiseq.extractors.vcf_matching import variants_to_pyranges, \
-    pyranges_to_intervals,  intervals_to_pyranges, BaseVariantMatcher, \
+    pyranges_to_intervals, intervals_to_pyranges, BaseVariantMatcher, \
     SingleVariantMatcher, MultiVariantsMatcher
 
 intervals = [
@@ -112,8 +112,7 @@ def test_BaseVariantMatcher__read_intervals():
 def test_SingleVariantMatcher__iter__():
     inters = intervals + [Interval('chr1', 5, 50)]
 
-    matcher = SingleVariantMatcher(
-        vcf_file, intervals=inters)
+    matcher = SingleVariantMatcher(vcf_file, intervals=inters)
     pairs = list(matcher)
 
     assert (inters[0], variants[0]) in pairs
@@ -134,6 +133,15 @@ def test_SingleVariantMatcher__iter__():
     assert (inters[2], variants[2]) in pairs
     assert len(pairs) == 5
 
+    matcher = SingleVariantMatcher(variants=variants, pranges=pr)
+    pairs = list(matcher)
+
+    assert (inters[0], variants[0]) in pairs
+    assert (inters[0], variants[1]) in pairs
+    assert (inters[1], variants[2]) in pairs
+    assert (inters[2], variants[1]) in pairs
+    assert (inters[2], variants[2]) in pairs
+    assert len(pairs) == 5
 
 def test_MultiVariantMatcher__iter__():
     matcher = MultiVariantsMatcher(vcf_file, intervals=intervals)
@@ -151,3 +159,11 @@ def test_MultiVariantMatcher__iter__():
     assert list(pairs[0][1]) == [variants[0], variants[1]]
     assert pairs[1][0] == intervals[1]
     assert list(pairs[1][1]) == [variants[2]]
+
+    matcher = MultiVariantsMatcher(variants=variants, pranges=pr)
+    pairs = list(matcher)
+
+    assert pairs[0][0] == intervals[0]
+    assert list(pairs[0][1]) == [variants[0], variants[1]]
+    assert pairs[1][0] == intervals[1]
+    assert list(pairs[1][1]) == [variants[2]]