Use individuals_nodes in vcf code

Author: benjeffery

python/CHANGELOG.rst

@@ -2,6 +2,15 @@
 [0.6.4] - 2025-XX-XX
 --------------------
 
+**Changes**
+
+- Allow mixed sample and non-sample nodes to be output by ``write_vcf``.
+  and for individuals without nodes to be specified for output, although
+  they will not appear in the VCF. Individuals without nodes no longer
+  raise an error and are instead ignored, except in the case where they
+  are explicitly specified in the ``individuals`` argument.
+  (:user:`benjeffery`, :pr:`3157`)
+
 **Features**
 
 - Add ``TreeSequence.sample_nodes_by_ploidy`` method to return the sample nodes

python/tests/test_vcf.py

@@ -340,7 +340,8 @@ def test_bad_ploidy(self):
         # 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)
 
@@ -358,8 +359,10 @@ def test_individuals_no_nodes_as_argument(self):
         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 pytest.raises(
+            ValueError, match="No samples were associated with the individuals"
+        ):
+            ts2.as_vcf(individuals=[0], allow_position_zero=True)
 
     def test_ploidy_with_sample_individuals(self):
         ts = msprime.sim_ancestry(3, random_seed=2)
@@ -397,10 +400,13 @@ def test_mixed_sample_non_sample_individuals(self):
         tables.nodes.individual = individual
         ts = tables.tree_sequence()
         ts = tsutil.insert_branch_sites(ts)
+        ts.as_vcf(allow_position_zero=True, isolated_as_missing=False)
         with pytest.raises(
-            ValueError, match="0 has nodes that are sample and non-sample"
+            tskit.LibraryError,
+            match="Cannot generate genotypes for non-samples when isolated "
+            "nodes are considered as missing",
         ):
-            ts.as_vcf()
+            ts.as_vcf(allow_position_zero=True)
         # but it's OK if we run without the affected individual
         assert len(ts.as_vcf(individuals=[1, 2], allow_position_zero=True)) > 0
 
@@ -414,12 +420,7 @@ 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
+        ts.as_vcf(allow_position_zero=True)
 
     def test_bad_individuals(self):
         ts = msprime.simulate(10, mutation_rate=0.1, random_seed=2)

python/tskit/trees.py

@@ -6432,14 +6432,9 @@ def write_vcf(
         with sample nodes are included in increasing order of individual ID.
 
         Subsets or permutations of the sample individuals may be specified
-        using the ``individuals`` argument. It is an error to specify any
-        individuals that are not associated with any nodes, or whose
-        nodes are not all samples.
-
-        Mixed-sample individuals (e.g., those associated with one node
-        that is a sample and another that is not) in the data model will
-        result in an error by default. However, such individuals can be
-        excluded using the ``individuals`` argument.
+        using the ``individuals`` argument. An individual specified that
+        is not associated with any nodes will be ignored and not included in the
+        output.
 
         If there are no individuals in the tree sequence,
         synthetic individuals are created by combining adjacent samples, and
@@ -10503,6 +10498,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 < 1:
+            raise ValueError("Ploidy must be >= 1")
         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
 
 
@@ -140,54 +139,32 @@ def __make_sample_mapping(self, ploidy, individuals):
             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"
-                    )
+        # If there are no individuals, or all the individuals are not associated with
+        # nodes, then we split by ploidy.
+        if ts.num_individuals > 0 and np.any(ts.individuals_nodes != -1):
+            individuals_nodes = ts.individuals_nodes
         else:
+            if ploidy is None:
+                ploidy = 1
+            individuals_nodes = ts.sample_nodes_by_ploidy(ploidy)
+
+        if individuals is not None:
             individuals = np.array(individuals, dtype=np.int32)
             if len(individuals) == 0:
                 raise ValueError("List of sample individuals empty")
+            if any(individuals < 0) or any(individuals >= ts.num_individuals):
+                raise ValueError("Invalid individual IDs provided.")
+            individuals_nodes = ts.individuals_nodes[individuals]
 
-        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.samples = []
+        for ind_nodes in individuals_nodes:
+            wanted_nodes = ind_nodes[ind_nodes != -1]
+            if len(wanted_nodes) > 0:
+                self.samples.extend(wanted_nodes)
+                self.individual_ploidies.append(len(wanted_nodes))
         self.num_individuals = len(self.individual_ploidies)
+        if self.num_individuals == 0:
+            raise ValueError("No samples were associated with the individuals")
 
     def __write_header(self, output):
         print("##fileformat=VCFv4.2", file=output)