Merge pull request #24 from RacimoLab/vcf

more vcf updates

Author: grahamgower

dinf/vcf.py

@@ -1,29 +1,16 @@
 from __future__ import annotations
 from typing import Dict, Iterable, List, Tuple
 import collections
-import contextlib
 import logging
-import os
 import pathlib
-import sys
+import warnings
 
 import numpy as np
 import cyvcf2
 
 logger = logging.getLogger(__name__)
 
 
-@contextlib.contextmanager
-def redirect_stderr(fp):
-    """Redirect stderr to fp."""
-    olderr = sys.stderr
-    try:
-        sys.stderr = fp
-        yield
-    finally:
-        sys.stderr = olderr
-
-
 def get_genotype_matrix(
     vcf: cyvcf2.VCF,
     *,
@@ -123,7 +110,8 @@ def __init__(
         self,
         files: Iterable[str | pathlib.Path],
         *,
-        individuals: List[str] | None = None,
+        individuals: Iterable[str] | None = None,
+        contigs: Iterable[str] | None = None,
     ):
         """
         :param files:
@@ -132,17 +120,34 @@ def __init__(
         :param individuals:
             An iterable of individual names corresponding to the VCF columns
             for which genotypes will be sampled.
+        :param contigs:
+            An iterable of contig names to use. If not specified, all contigs
+            identified in the VCFs/BCFs will be used.
         """
-        with open(os.devnull, "w") as dev_null:
-            self._fill_bag(files=files, individuals=individuals, dev_null=dev_null)
+        # Silence some warnings from cyvcf2.
+        with warnings.catch_warnings():
+            # We check if a contig is usable for a given vcf by querying
+            # that contig for variants. If there are no variants (e.g. vcfs are
+            # split by chromosome so each vcf has data for only one contig),
+            # then cyvcf2 warns us.
+            warnings.filterwarnings(
+                "ignore", message="no intervals found", category=UserWarning
+            )
+            # We detect this condition and explicitly raise an error.
+            warnings.filterwarnings(
+                "ignore",
+                message="not all requested samples found",
+                category=UserWarning,
+            )
+            self._fill_bag(files=files, individuals=individuals, contigs=contigs)
         self._regions: List[Tuple[str, int, int]] = []
 
     def _fill_bag(
         self,
         *,
         files: Iterable[str | pathlib.Path],
-        individuals: List[str] | None = None,
-        dev_null,
+        individuals: Iterable[str] | None = None,
+        contigs: Iterable[str] | None = None,
     ) -> None:
         """
         Construct a mapping from contig id to cyvcf2.VCF object.
@@ -156,6 +161,19 @@ def _fill_bag(
         if len(set(files)) != len(files):
             raise ValueError("File list contains duplicates.")
 
+        if individuals is not None:
+            individuals = list(individuals)
+            if len(set(individuals)) != len(individuals):
+                raise ValueError("Individuals list contains duplicates.")
+
+        if contigs is not None:
+            contigs = list(contigs)
+            if len(set(contigs)) != len(contigs):
+                raise ValueError("Contigs list contains duplicates.")
+            contigs = set(contigs)
+
+        contigs_seen = set()
+
         for file in files:
             vcf = cyvcf2.VCF(file, samples=individuals, lazy=True, threads=1)
             if "GT" not in vcf:
@@ -165,15 +183,23 @@ def _fill_bag(
                 or pathlib.Path(f"{file}.csi").exists()
             ):
                 raise ValueError(f"No index found for {file}.")
+            if individuals is not None:
+                individuals_not_found = set(individuals) - set(vcf.samples)
+                if len(individuals_not_found) > 0:
+                    raise ValueError(
+                        f"Requested individuals not found in {file}: "
+                        f"{', '.join(individuals_not_found)}"
+                    )
             for contig_id, contig_length in zip(vcf.seqnames, vcf.seqlens):
+                contigs_seen.add(contig_id)
+                if contigs is not None and contig_id not in contigs:
+                    continue
                 # Check there's at least one variant. If present, we assume
                 # there exist usable variants (i.e. SNPs) for the contig.
-                # We redirect stderr to silence cyvcf2 "no intervals" warnings.
-                with redirect_stderr(dev_null):
-                    try:
-                        next(vcf(contig_id))
-                    except StopIteration:
-                        continue
+                try:
+                    next(vcf(contig_id))
+                except StopIteration:
+                    continue
                 if contig_id in contig2file:
                     first_file = contig2file[contig_id]
                     raise ValueError(
@@ -184,6 +210,13 @@ def _fill_bag(
                 contig_lengths.append(contig_length)
                 contig2vcf[contig_id] = vcf
 
+        if contigs is not None:
+            contigs_not_found = contigs - contigs_seen
+            if len(contigs_not_found) > 0:
+                raise ValueError(
+                    f"Requested contigs not found: {', '.join(contigs_not_found)}"
+                )
+
         if len(contig2file) == 0:
             raise ValueError("No usable vcf/bcf files in the list.")
 

environment-gpu.yml

@@ -17,10 +17,10 @@ dependencies:
   - optax
   # The rest.
   - arviz
-  - cyvcf2
+  - cyvcf2 >= 0.30.14
   - emcee
-  - demes
-  - msprime
+  - demes >= 0.2.1
+  - msprime >= 1.0.4
   - pip
   - pip:
     # Last jaxlib version to support CUDA 10.1

environment.yml

@@ -13,7 +13,7 @@ dependencies:
   - optax
   # The rest.
   - arviz
-  - cyvcf2
+  - cyvcf2 >= 0.30.14
   - emcee
-  - demes
-  - msprime
+  - demes >= 0.2.1
+  - msprime >= 1.0.4

requirements/minimal.txt

@@ -1,10 +1,10 @@
 arviz==0.11.4
-cyvcf2==0.30.12
+cyvcf2==0.30.14
 emcee==3.1.1
 numpy==1.21.4
 jax==0.2.25
 jaxlib==0.1.74
 flax==0.3.6
 optax==0.1.0
-msprime==1.0.3
-demes==0.1.2
+msprime==1.0.4
+demes==0.2.1

setup.cfg

@@ -30,14 +30,14 @@ include_package_data = True
 python_requires = >=3.7
 install_requires =
     arviz
-    cyvcf2
+    cyvcf2 >= 0.30.14
     emcee
     numpy
     jax
     flax
     optax
     msprime >= 1.0.4
-    demes
+    demes >= 0.2.1
 setup_requires =
     setuptools
     setuptools_scm

tests/test_vcf.py

@@ -487,8 +487,17 @@ def test_unindexed_bcf(self, tmp_path):
         with pytest.raises(ValueError, match="No index"):
             dinf.BagOfVcf(tmp_path.glob("*.bcf"))
 
+    def test_missing_file(self):
+        missing_file = "nonexistent.vcf.gz"
+        with pytest.raises(OSError, match=missing_file):
+            dinf.BagOfVcf([missing_file])
+
+    def test_no_files(self):
+        with pytest.raises(ValueError, match="No usable vcf/bcf files"):
+            dinf.BagOfVcf([])
+
     @pytest.mark.usefixtures("tmp_path")
-    def test_file_list_duplicates(self, tmp_path):
+    def test_duplicate_files(self, tmp_path):
         create_vcf_dataset(tmp_path, contig_lengths=[100_000])
         files = 2 * list(tmp_path.glob("*.vcf.gz"))
         with pytest.raises(ValueError, match="File list contains duplicates"):
@@ -504,10 +513,6 @@ def test_multiple_files_claim_a_contig(self, tmp_path):
         with pytest.raises(ValueError, match="Both .* contain records for sequence"):
             dinf.BagOfVcf(files)
 
-    def test_no_files(self):
-        with pytest.raises(ValueError, match="No usable vcf/bcf files"):
-            dinf.BagOfVcf([])
-
     @pytest.mark.usefixtures("tmp_path")
     def test_no_GT_field(self, tmp_path):
         def remove_GT_header(filename):
@@ -523,6 +528,57 @@ def remove_GT_header(filename):
         with pytest.raises(ValueError, match="doesn't contain GT field"):
             dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"))
 
+    @pytest.mark.parametrize("num_individuals", [20, 97])
+    @pytest.mark.usefixtures("tmp_path")
+    def test_individuals(self, tmp_path, num_individuals):
+        samples = create_vcf_dataset(tmp_path, contig_lengths=[100_000])
+        individuals = samples["A"][:num_individuals]
+        vb = dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"), individuals=individuals)
+        assert vb["1"].samples == individuals
+
+    @pytest.mark.filterwarnings("ignore:not all requested samples found:UserWarning")
+    @pytest.mark.usefixtures("tmp_path")
+    def test_missing_individuals(self, tmp_path):
+        samples = create_vcf_dataset(tmp_path, contig_lengths=[100_000])
+        individuals = ["nonexistent_1"] + samples["A"] + ["nonexistent_2"]
+        with pytest.raises(ValueError, match="individuals not found") as err:
+            dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"), individuals=individuals)
+        assert "nonexistent_1" in err.value.args[0]
+        assert "nonexistent_2" in err.value.args[0]
+        for ind in samples["A"]:
+            assert ind not in err.value.args[0]
+
+    @pytest.mark.usefixtures("tmp_path")
+    def test_duplicate_individuals(self, tmp_path):
+        samples = create_vcf_dataset(tmp_path, contig_lengths=[100_000])
+        individuals = samples["A"] + [samples["A"][0]]
+        with pytest.raises(ValueError, match="Individuals list contains duplicates"):
+            dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"), individuals=individuals)
+
+    @pytest.mark.usefixtures("tmp_path")
+    def test_contigs(self, tmp_path):
+        create_vcf_dataset(tmp_path, contig_lengths=[100_000, 200_000, 300_000])
+        contigs = ["1", "3"]
+        vb = dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"), contigs=contigs)
+        assert set(vb) == set(contigs)
+
+    @pytest.mark.usefixtures("tmp_path")
+    def test_missing_contigs(self, tmp_path):
+        create_vcf_dataset(tmp_path, contig_lengths=[100_000, 200_000])
+        contigs = ["nonexistent_a", "1", "2", "nonexistent_b"]
+        with pytest.raises(ValueError, match="contigs not found") as err:
+            dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"), contigs=contigs)
+        assert "nonexistent_a" in err.value.args[0]
+        assert "nonexistent_b" in err.value.args[0]
+        assert "1" not in err.value.args[0]
+        assert "2" not in err.value.args[0]
+
+    @pytest.mark.usefixtures("tmp_path")
+    def test_duplicate_contigs(self, tmp_path):
+        create_vcf_dataset(tmp_path, contig_lengths=[100_000])
+        with pytest.raises(ValueError, match="Contigs list contains duplicates"):
+            dinf.BagOfVcf(tmp_path.glob("*.vcf.gz"), contigs=["1", "1"])
+
     @pytest.mark.usefixtures("tmp_path")
     def test_unused_contigs(self, tmp_path):
         # Contigs in the header should be ignored if they have no variants.