Use vcf model in vcf code

Author: benjeffery

python/tests/test_vcf.py

@@ -29,6 +29,7 @@
 import os
 import tempfile
 import textwrap
+import warnings
 
 import msprime
 import numpy as np
@@ -335,12 +336,13 @@ class TestInterface:
     def test_bad_ploidy(self):
         ts = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
         for bad_ploidy in [-1, 0]:
-            with pytest.raises(ValueError, match="Ploidy must be >= 1"):
+            with pytest.raises(ValueError, match="Ploidy must be a positive integer"):
                 ts.write_vcf(io.StringIO, bad_ploidy)
         # Non divisible
         for bad_ploidy in [3, 7]:
             with pytest.raises(
-                ValueError, match="Sample size must be divisible by ploidy"
+                ValueError,
+                match="Number of sample nodes 10 is not a multiple of ploidy",
             ):
                 ts.write_vcf(io.StringIO, bad_ploidy)
 
@@ -349,17 +351,32 @@ def test_individuals_no_nodes_default_args(self):
         tables = ts1.dump_tables()
         tables.individuals.add_row()
         ts2 = tables.tree_sequence()
-        assert ts1.as_vcf(allow_position_zero=True) == ts2.as_vcf(
-            allow_position_zero=True
-        )
+        # ts1 should work as it has no individuals
+        ts1.as_vcf(allow_position_zero=True)
+        # ts2 should fail as it has individuals but no nodes
+        with warnings.catch_warnings(record=True) as w:
+            with pytest.raises(ValueError, match="No samples in resulting VCF model"):
+                ts2.as_vcf(allow_position_zero=True)
+            assert len(w) == 2
+            assert "At least one sample node does not have an individual ID" in str(
+                w[0].message
+            )
+            assert "Individual 0 has no nodes associated with it." in str(w[1].message)
 
     def test_individuals_no_nodes_as_argument(self):
         ts1 = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
         tables = ts1.dump_tables()
         tables.individuals.add_row()
         ts2 = tables.tree_sequence()
-        with pytest.raises(ValueError, match="0 not associated with a node"):
-            ts2.as_vcf(individuals=[0])
+        with warnings.catch_warnings(record=True) as w:
+            with pytest.raises(
+                ValueError, match="Individual 0 has no nodes associated with it."
+            ):
+                ts2.as_vcf(individuals=[0])
+            assert len(w) == 1
+            assert "At least one sample node does not have an individual ID" in str(
+                w[0].message
+            )
 
     def test_ploidy_with_sample_individuals(self):
         ts = msprime.sim_ancestry(3, random_seed=2)
@@ -378,7 +395,7 @@ def test_ploidy_with_no_node_individuals(self):
     def test_empty_individuals(self):
         ts = msprime.sim_ancestry(3, random_seed=2)
         ts = tsutil.insert_branch_sites(ts)
-        with pytest.raises(ValueError, match="List of sample individuals empty"):
+        with pytest.raises(ValueError, match="No samples in resulting VCF model"):
             ts.as_vcf(individuals=[])
 
     def test_duplicate_individuals(self):
@@ -397,12 +414,17 @@ def test_mixed_sample_non_sample_individuals(self):
         tables.nodes.individual = individual
         ts = tables.tree_sequence()
         ts = tsutil.insert_branch_sites(ts)
-        with pytest.raises(
-            ValueError, match="0 has nodes that are sample and non-sample"
-        ):
-            ts.as_vcf()
-        # but it's OK if we run without the affected individual
-        assert len(ts.as_vcf(individuals=[1, 2], allow_position_zero=True)) > 0
+        with warnings.catch_warnings(record=True) as w:
+            ts.map_to_vcf_model()
+            assert len(w) == 2
+            assert (
+                "Individual 0 has both sample and non-sample nodes associated with it."
+                in str(w[0].message)
+            )
+            assert "Individual 3 has no nodes associated with it." in str(w[1].message)
+        with warnings.catch_warnings(record=True) as w:
+            assert len(ts.as_vcf(individuals=[1, 2], allow_position_zero=True)) > 0
+            assert len(w) == 0
 
     def test_samples_with_and_without_individuals(self):
         ts = tskit.Tree.generate_balanced(3).tree_sequence
@@ -414,19 +436,19 @@ def test_samples_with_and_without_individuals(self):
         tables.nodes.individual = individual
         ts = tables.tree_sequence()
         ts = tsutil.insert_branch_sites(ts)
-        with pytest.raises(
-            ValueError, match="Sample nodes must either all be associated"
-        ):
-            ts.as_vcf()
-        # But it's OK if explicitly specify that sample
-        assert len(ts.as_vcf(individuals=[0], allow_position_zero=True)) > 0
+        with warnings.catch_warnings(record=True) as w:
+            ts.as_vcf(allow_position_zero=True)
+            assert len(w) == 1
+            assert "At least one sample node does not have an individual ID" in str(
+                w[0].message
+            )
 
     def test_bad_individuals(self):
         ts = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
         ts = tsutil.insert_individuals(ts, ploidy=2)
-        with pytest.raises(ValueError, match="Invalid individual IDs provided."):
+        with pytest.raises(ValueError, match="Invalid individual ID"):
             ts.write_vcf(io.StringIO(), individuals=[0, -1])
-        with pytest.raises(ValueError, match="Invalid individual IDs provided."):
+        with pytest.raises(ValueError, match="Invalid individual ID"):
             ts.write_vcf(io.StringIO(), individuals=[1, 2, ts.num_individuals])
 
     def test_ploidy_positional(self):
@@ -527,20 +549,17 @@ def test_bad_length_individuals(self):
         ts = tsutil.insert_individuals(ts, ploidy=2)
         with pytest.raises(
             ValueError,
-            match="individual_names must have length equal to"
-            " the number of individuals",
+            match="The number of individuals does not match the number of names",
         ):
             ts.write_vcf(io.StringIO(), individual_names=[])
         with pytest.raises(
             ValueError,
-            match="individual_names must have length equal to"
-            " the number of individuals",
+            match="The number of individuals does not match the number of names",
         ):
             ts.write_vcf(io.StringIO(), individual_names=["x" for _ in range(4)])
         with pytest.raises(
             ValueError,
-            match="individual_names must have length equal to"
-            " the number of individuals",
+            match="The number of individuals does not match the number of names",
         ):
             ts.write_vcf(
                 io.StringIO(),
@@ -549,8 +568,7 @@ def test_bad_length_individuals(self):
             )
         with pytest.raises(
             ValueError,
-            match="individual_names must have length equal to"
-            " the number of individuals",
+            match="The number of individuals does not match the number of names",
         ):
             ts.write_vcf(
                 io.StringIO(),
@@ -563,14 +581,12 @@ def test_bad_length_ploidy(self):
         assert ts.num_sites > 0
         with pytest.raises(
             ValueError,
-            match="individual_names must have length equal to"
-            " the number of individuals",
+            match="The number of individuals does not match the number of names",
         ):
             ts.write_vcf(io.StringIO(), ploidy=2, individual_names=[])
         with pytest.raises(
             ValueError,
-            match="individual_names must have length equal to"
-            " the number of individuals",
+            match="The number of individuals does not match the number of names",
         ):
             ts.write_vcf(
                 io.StringIO(), ploidy=2, individual_names=["x" for _ in range(4)]
@@ -907,7 +923,7 @@ def test_no_individuals_ploidy_3_names(self):
         expected = textwrap.dedent(s)
         assert (
             drop_header(
-                ts.as_vcf(ploidy=3, individual_names="A", allow_position_zero=True)
+                ts.as_vcf(ploidy=3, individual_names=["A"], allow_position_zero=True)
             )
             == expected
         )
@@ -940,7 +956,7 @@ def test_individual_0(self):
     def test_individual_1(self):
         ts = self.ts()
         s = """\
-        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0
+        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_1
         1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t0
         1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t1
         1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t1
@@ -954,7 +970,7 @@ def test_individual_1(self):
     def test_reversed(self):
         ts = self.ts()
         s = """\
-        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1
+        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_1\ttsk_0
         1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t0\t1|0
         1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t1\t0|1
         1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t1\t0|0

python/tskit/trees.py

@@ -10513,6 +10513,8 @@ def sample_nodes_by_ploidy(self, ploidy):
         :return: A 2D array of node IDs, where each row has length `ploidy`.
         :rtype: numpy.ndarray
         """
+        if ploidy <= 0 or ploidy != int(ploidy):
+            raise ValueError("Ploidy must be a positive integer")
         sample_node_ids = np.flatnonzero(self.nodes_flags & tskit.NODE_IS_SAMPLE)
         num_samples = len(sample_node_ids)
         if num_samples == 0:

python/tskit/vcf.py

@@ -25,7 +25,6 @@
 """
 import numpy as np
 
-import tskit
 from . import provenance
 
 
@@ -69,14 +68,24 @@ def __init__(
         self.contig_id = contig_id
         self.isolated_as_missing = isolated_as_missing
 
-        self.__make_sample_mapping(ploidy, individuals)
-        if individual_names is None:
-            individual_names = [f"tsk_{j}" for j in range(self.num_individuals)]
-        self.individual_names = individual_names
-        if len(self.individual_names) != self.num_individuals:
-            raise ValueError(
-                "individual_names must have length equal to the number of individuals"
-            )
+        vcf_model = tree_sequence.map_to_vcf_model(
+            individuals=individuals, ploidy=ploidy, individual_names=individual_names
+        )
+        self.individual_names = vcf_model.individuals_name
+        self.individual_ploidies = [
+            len(nodes[nodes >= 0]) for nodes in vcf_model.individuals_nodes
+        ]
+        self.num_individuals = len(self.individual_names)
+
+        if len(vcf_model.individuals_nodes) == 0:
+            raise ValueError("No samples in resulting VCF model")
+
+        # Flatten the array of node IDs, filtering out the -1 padding values
+        self.samples = []
+        for row in vcf_model.individuals_nodes:
+            for node_id in row:
+                if node_id != -1:
+                    self.samples.append(node_id)
 
         # Transform coordinates for VCF
         if position_transform is None:
@@ -126,68 +135,7 @@ def __init__(
                 '"position_transform = lambda x: np.fmax(1, x)"'
             )
 
-    def __make_sample_mapping(self, ploidy, individuals):
-        """
-        Compute the sample IDs for each VCF individual and the template for
-        writing out genotypes.
-        """
-        ts = self.tree_sequence
-        self.samples = None
-        self.individual_ploidies = []
-
-        # Cannot use "ploidy" when *any* individuals are present.
-        if ts.num_individuals > 0 and ploidy is not None:
-            raise ValueError(
-                "Cannot specify ploidy when individuals are present in tables "
-            )
-
-        if individuals is None:
-            # Find all sample nodes that reference individuals
-            individuals = np.unique(ts.nodes_individual[ts.samples()])
-            if len(individuals) == 1 and individuals[0] == tskit.NULL:
-                # No samples refer to individuals
-                individuals = None
-            else:
-                # np.unique sorts the argument, so if NULL (-1) is present it
-                # will be the first value.
-                if individuals[0] == tskit.NULL:
-                    raise ValueError(
-                        "Sample nodes must either all be associated with individuals "
-                        "or not associated with any individuals"
-                    )
-        else:
-            individuals = np.array(individuals, dtype=np.int32)
-            if len(individuals) == 0:
-                raise ValueError("List of sample individuals empty")
 
-        if individuals is not None:
-            self.samples = []
-            # FIXME this could probably be done more efficiently.
-            for i in individuals:
-                if i < 0 or i >= self.tree_sequence.num_individuals:
-                    raise ValueError("Invalid individual IDs provided.")
-                ind = self.tree_sequence.individual(i)
-                if len(ind.nodes) == 0:
-                    raise ValueError(f"Individual {i} not associated with a node")
-                is_sample = {ts.node(u).is_sample() for u in ind.nodes}
-                if len(is_sample) != 1:
-                    raise ValueError(
-                        f"Individual {ind.id} has nodes that are sample and "
-                        "non-samples"
-                    )
-                self.samples.extend(ind.nodes)
-                self.individual_ploidies.append(len(ind.nodes))
-        else:
-            if ploidy is None:
-                ploidy = 1
-            if ploidy < 1:
-                raise ValueError("Ploidy must be >= 1")
-            if ts.num_samples % ploidy != 0:
-                raise ValueError("Sample size must be divisible by ploidy")
-            self.individual_ploidies = np.full(
-                ts.sample_size // ploidy, ploidy, dtype=np.int32
-            )
-        self.num_individuals = len(self.individual_ploidies)
 
     def __write_header(self, output):
         print("##fileformat=VCFv4.2", file=output)