Support for specifying a list of contigs to use.

Author: grahamgower

dinf/vcf.py

@@ -111,6 +111,7 @@ def __init__(
         files: Iterable[str | pathlib.Path],
         *,
         individuals: Iterable[str] | None = None,
+        contigs: Iterable[str] | None = None,
     ):
         """
         :param files:
@@ -119,6 +120,9 @@ 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.
         """
         # Silence some warnings from cyvcf2.
         with warnings.catch_warnings():
@@ -135,14 +139,15 @@ def __init__(
                 message="not all requested samples found",
                 category=UserWarning,
             )
-            self._fill_bag(files=files, individuals=individuals)
+            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: Iterable[str] | None = None,
+        contigs: Iterable[str] | None = None,
     ) -> None:
         """
         Construct a mapping from contig id to cyvcf2.VCF object.
@@ -155,8 +160,19 @@ def _fill_bag(
         files = list(files)
         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)
@@ -168,13 +184,16 @@ def _fill_bag(
             ):
                 raise ValueError(f"No index found for {file}.")
             if individuals is not None:
-                not_found = set(individuals) - set(vcf.samples)
-                if len(not_found) > 0:
+                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(not_found)}"
+                        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.
                 try:
@@ -191,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.")
 

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):
@@ -533,7 +538,7 @@ def test_individuals(self, tmp_path, num_individuals):
 
     @pytest.mark.filterwarnings("ignore:not all requested samples found:UserWarning")
     @pytest.mark.usefixtures("tmp_path")
-    def test_bad_individuals(self, 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:
@@ -543,6 +548,37 @@ def test_bad_individuals(self, tmp_path):
         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.