Merge pull request #100 from Hoeze/generalize_variant_source

Allow to directly annotate a list of variants

Author: Hoeze

.circleci/config.yml

@@ -117,7 +117,7 @@ variables:
 jobs:
   test-py36:
     docker:
-      - image: continuumio/miniconda3:4.7.10
+      - image: continuumio/miniconda3:4.10.3
     working_directory: ~/repo
     steps:
       - checkout
@@ -135,7 +135,7 @@ jobs:
 
   build-deploy-docs:
     docker:
-      - image: continuumio/miniconda3:4.7.10
+      - image: continuumio/miniconda3:4.10.3
     working_directory: ~/repo
     steps:
       - add_ssh_keys:
@@ -168,7 +168,7 @@ jobs:
 
   test-deploy-pypi:
     docker:
-      - image: continuumio/miniconda3:4.7.10
+      - image: continuumio/miniconda3:4.10.3
     working_directory: ~/repo
     steps:
       - checkout
@@ -197,7 +197,7 @@ jobs:
 
   productive-deploy-pypi:
     docker:
-      - image: continuumio/miniconda3:4.7.10
+      - image: continuumio/miniconda3:4.10.3
     working_directory: ~/repo
     steps:
       - checkout

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,43 @@ 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 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 +131,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 +141,31 @@ 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,
+    ) -> VariantFetcher:
+        if vcf_file is not None:
+            from kipoiseq.extractors import MultiSampleVCF
+            return MultiSampleVCF(vcf_file, lazy=vcf_lazy)
+        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 +222,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 +260,14 @@ def __iter__(self) -> (Interval, Variant):
 
 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

setup.py

@@ -34,10 +34,10 @@
     "cython",
     "cyvcf2",
     "pyranges>=0.0.71",
-    "keras",
-    "tensorflow",
+    # "keras",
+    # "tensorflow",
     "pybedtools",
-    "concise"
+    # "concise"
 ]
 
 setup(

tests/extractors/test_vcf_matching.py

@@ -1,11 +1,12 @@
+from typing import Union, Iterable, Iterator, List
 import pytest
 from conftest import vcf_file, gtf_file, example_intervals_bed
 import pyranges
 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, \
-    SingleVariantMatcher, MultiVariantsMatcher
+    SingleVariantMatcher, MultiVariantsMatcher, VariantFetcher
 
 intervals = [
     Interval('chr1', 1, 10, strand='+'),
@@ -26,6 +27,36 @@
 )
 
 
+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
+
+
+# make sure that kipoiseq only uses the VariantFetcher API
+read_variants_fn = BaseVariantMatcher._read_variants
+
+
+@staticmethod
+def proxy_fn(*args, **kwargs):
+    vf = VariantFetcherProxy(
+        read_variants_fn(*args, **kwargs)
+    )
+    return vf
+
+
+BaseVariantMatcher._read_variants = proxy_fn
+
+
 def test_variants_to_pyranges():
     vcf = MultiSampleVCF(vcf_file)
     variants = list(vcf)
@@ -131,6 +162,15 @@ def test_SingleVariantMatcher__iter__():
     assert (inters[2], variants[2]) in pairs
     assert len(pairs) == 4
 
+    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[2]) in pairs
+    assert len(pairs) == 4
+
 
 def test_MultiVariantMatcher__iter__():
     matcher = MultiVariantsMatcher(vcf_file, intervals=intervals)
@@ -148,3 +188,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]]