Add additional tests for vcf mapping

Author: benjeffery

python/tests/test_vcf.py

@@ -975,3 +975,78 @@ def test_reversed_names(self):
             )
             == expected
         )
+
+
+class TestVcfMapping:
+    def test_mix_sample_non_sample(self):
+        ts = tskit.Tree.generate_balanced(5, span=10).tree_sequence
+        ts = tsutil.insert_branch_sites(ts)
+        assert ts.num_nodes >= 8
+        tables = ts.dump_tables()
+        tables.individuals.add_row()
+        tables.individuals.add_row()
+        tables.individuals.add_row()
+        tables.individuals.add_row()
+        individual = tables.nodes.individual
+        assert np.all(individual == -1)
+        # First has only non-sample nodes
+        individual[7] = 0
+        # Second has 2 sample nodes
+        individual[0] = 1
+        individual[1] = 1
+        # Third has 1 non-sample and 1 sample
+        individual[5] = 2
+        individual[2] = 2
+        # Fourth has sandwiched non-sample
+        individual[3] = 3
+        individual[6] = 3
+        individual[4] = 3
+        tables.nodes.individual = individual
+        ts = tables.tree_sequence()
+
+        # Individual "A" is redacted as has no nodes
+        s = """\
+        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tB\tC\tD
+        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1|1\t0\t0|0
+        1\t1\t1\t0\t1\t.\tPASS\t.\tGT\t1|0\t0\t0|0
+        1\t2\t2\t0\t1\t.\tPASS\t.\tGT\t0|1\t0\t0|0
+        1\t3\t3\t0\t1\t.\tPASS\t.\tGT\t0|0\t1\t1|1
+        1\t4\t4\t0\t1\t.\tPASS\t.\tGT\t0|0\t1\t0|0
+        1\t6\t5\t0\t1\t.\tPASS\t.\tGT\t0|0\t0\t1|1
+        1\t7\t6\t0\t1\t.\tPASS\t.\tGT\t0|0\t0\t1|0
+        1\t8\t7\t0\t1\t.\tPASS\t.\tGT\t0|0\t0\t0|1"""
+        expected = textwrap.dedent(s)
+        assert (
+            drop_header(
+                ts.as_vcf(
+                    individual_names=["A", "B", "C", "D"],
+                    allow_position_zero=True,
+                ),
+            )
+            == expected
+        )
+
+        # Now with non-sample nodes, so A is included, C becomes diploid
+        # and D is triploid
+        s = """\
+        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tA\tB\tC\tD
+        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t0\t1|1\t0|1\t0|0|0
+        1\t1\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1|0\t0|0\t0|0|0
+        1\t2\t2\t0\t1\t.\tPASS\t.\tGT\t0\t0|1\t0|0\t0|0|0
+        1\t3\t3\t0\t1\t.\tPASS\t.\tGT\t1\t0|0\t1|0\t1|1|1
+        1\t4\t4\t0\t1\t.\tPASS\t.\tGT\t0\t0|0\t1|0\t0|0|0
+        1\t6\t5\t0\t1\t.\tPASS\t.\tGT\t0\t0|0\t0|0\t1|1|1
+        1\t7\t6\t0\t1\t.\tPASS\t.\tGT\t0\t0|0\t0|0\t1|0|0
+        1\t8\t7\t0\t1\t.\tPASS\t.\tGT\t0\t0|0\t0|0\t0|1|0"""
+        expected = textwrap.dedent(s)
+        assert (
+            drop_header(
+                ts.as_vcf(
+                    individual_names=["A", "B", "C", "D"],
+                    allow_position_zero=True,
+                    include_non_sample_nodes=True,
+                    isolated_as_missing=False,
+                ),
+            )
+            == expected
+        )

python/tskit/trees.py

@@ -6416,6 +6416,7 @@ def write_vcf(
         sample_mask=None,
         isolated_as_missing=None,
         allow_position_zero=None,
+        include_non_sample_nodes=None,
     ):
         """
         Convert the genetic variation data in this tree sequence to Variant
@@ -6559,6 +6560,7 @@ def write_vcf(
             sample_mask=sample_mask,
             isolated_as_missing=isolated_as_missing,
             allow_position_zero=allow_position_zero,
+            include_non_sample_nodes=include_non_sample_nodes,
         )
         writer.write(output)
 

python/tskit/vcf.py

@@ -63,23 +63,26 @@ def __init__(
         sample_mask,
         isolated_as_missing,
         allow_position_zero,
+        include_non_sample_nodes,
     ):
         self.tree_sequence = tree_sequence
         self.contig_id = contig_id
         self.isolated_as_missing = isolated_as_missing
 
         vcf_model = tree_sequence.map_to_vcf_model(
-            individuals=individuals, ploidy=ploidy, individual_names=individual_names
+            individuals=individuals,
+            ploidy=ploidy,
+            individual_names=individual_names,
+            include_non_sample_nodes=include_non_sample_nodes,
         )
         # Remove individuals with zero ploidy as these cannot be
         # represented in VCF.
         individuals_nodes = vcf_model.individuals_nodes
         to_keep = (individuals_nodes != -1).any(axis=1)
         individuals_nodes = individuals_nodes[to_keep]
         self.individual_names = vcf_model.individuals_name[to_keep]
-
         self.individual_ploidies = [
-            len(nodes[nodes >= 0]) for nodes in vcf_model.individuals_nodes
+            len(nodes[nodes >= 0]) for nodes in individuals_nodes
         ]
         self.num_individuals = len(self.individual_names)