Fix map_to_vcf_model tests

Author: benjeffery

python/tests/test_highlevel.py

@@ -5747,23 +5747,15 @@ def test_specific_individuals(self):
         result = ts.map_to_vcf_model(individuals=[1, 3])
         assert isinstance(result, tskit.VcfModelMapping)
         # Individual 1 has ploidy 2, individual 3 has ploidy 4
-        assert result.individuals_nodes.shape == (2, 4)
-
-        assert result.individuals_nodes[0, 0] == 1
-        assert result.individuals_nodes[0, 1] == 2
-        assert result.individuals_nodes[0, 2] == -1
-        assert result.individuals_nodes[0, 3] == -1
-
-        assert result.individuals_nodes[1, 0] == 6
-        assert result.individuals_nodes[1, 1] == 7
-        assert result.individuals_nodes[1, 2] == 8
-        assert result.individuals_nodes[1, 3] == 9
+        assert result.individuals_nodes.shape == (2, 5)
+        assert np.array_equal(result.individuals_nodes[0], [1, 2, -1, -1, -1])
+        assert np.array_equal(result.individuals_nodes[1], [6, 7, 8, 9, -1])
 
         assert result.individuals_name.shape == (2,)
         assert result.individuals_name[0] == "tsk_1"
         assert result.individuals_name[1] == "tsk_3"
 
-    def test_individual_with_no_nodes_warning(self):
+    def test_individual_with_no_nodes(self):
         tables = tskit.TableCollection(1.0)
         # Individual with no nodes
         tables.individuals.add_row()
@@ -5772,18 +5764,11 @@ def test_individual_with_no_nodes_warning(self):
         tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
         ts = tables.tree_sequence()
 
-        with warnings.catch_warnings(record=True) as w:
-            result = ts.map_to_vcf_model()
-            assert len(w) == 1
-            assert "Individual 0 has no nodes" in str(w[0].message)
-
-        # Should only include individual 1
-        assert result.individuals_nodes.shape == (1, 1)
-        assert result.individuals_nodes[0, 0] == 0
-        assert result.individuals_name.shape == (1,)
-        assert result.individuals_name[0] == "tsk_1"
+        result = ts.map_to_vcf_model()
+        assert result.individuals_nodes.shape == (2, 1)
+        assert np.array_equal(result.individuals_nodes, [[-1], [0]])
 
-    def test_individual_with_no_nodes_error(self):
+    def test_individual_with_no_nodes_only(self):
         tables = tskit.TableCollection(1.0)
         # Individual with no nodes
         tables.individuals.add_row()
@@ -5792,8 +5777,9 @@ def test_individual_with_no_nodes_error(self):
         tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
         ts = tables.tree_sequence()
 
-        with pytest.raises(ValueError, match="Individual 0 has no nodes"):
-            ts.map_to_vcf_model(individuals=[0])
+        result = ts.map_to_vcf_model(individuals=[0])
+        assert result.individuals_nodes.shape == (1, 1)
+        assert np.array_equal(result.individuals_nodes, [[-1]])
 
     def test_invalid_individual_id(self):
         tables = tskit.TableCollection(1.0)
@@ -5807,19 +5793,31 @@ def test_invalid_individual_id(self):
         with pytest.raises(ValueError, match="Invalid individual ID"):
             ts.map_to_vcf_model(individuals=[1])
 
-    def test_mixed_sample_non_sample_warning(self):
+    def test_mixed_sample_non_sample_ordering(self):
         tables = tskit.TableCollection(1.0)
         tables.individuals.add_row()
         tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
         tables.nodes.add_row(flags=0, time=0, individual=0)  # Non-sample node
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=0, time=0, individual=0)  # Non-sample node
+        tables.individuals.add_row()
+        tables.nodes.add_row(flags=0, time=0, individual=1)  # Non-sample node
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
         ts = tables.tree_sequence()
 
-        with warnings.catch_warnings(record=True) as w:
-            ts.map_to_vcf_model()
-            assert len(w) == 1
-            assert "Individual 0 has both sample and non-sample nodes" in str(
-                w[0].message
-            )
+        result = ts.map_to_vcf_model()
+        assert result.individuals_nodes.shape == (2, 4)
+        assert np.array_equal(
+            result.individuals_nodes,
+            np.array([[0, 2, -1, -1], [5, -1, -1, -1]]),
+        )
+
+        result = ts.map_to_vcf_model(include_non_sample_nodes=True)
+        assert result.individuals_nodes.shape == (2, 4)
+        assert np.array_equal(
+            result.individuals_nodes,
+            np.array([[0, 1, 2, 3], [4, 5, -1, -1]]),
+        )
 
     def test_samples_without_individuals_warning(self):
         tables = tskit.TableCollection(1.0)
@@ -5898,3 +5896,11 @@ def test_name_count_mismatch_error(self):
             ValueError, match="number of individuals does not match the number of names"
         ):
             ts.map_to_vcf_model(individual_names=["only_one_name"])
+
+    def test_all_individuals_no_nodes(self):
+        tables = tskit.TableCollection(1.0)
+        tables.individuals.add_row()
+        tables.individuals.add_row()
+        ts = tables.tree_sequence()
+        result = ts.map_to_vcf_model()
+        assert result.individuals_nodes.shape == (2, 0)

python/tskit/trees.py

@@ -10608,12 +10608,17 @@ def map_to_vcf_model(
                 raise ValueError("Invalid individual ID")
 
             individuals_nodes = self.individuals_nodes[individuals]
-            if not include_non_sample_nodes:
-                individuals_nodes[
-                    np.logical_not(
-                        self.nodes_flags[individuals_nodes] & tskit.NODE_IS_SAMPLE
+            non_sample_nodes = np.logical_not(
+                self.nodes_flags[individuals_nodes] & tskit.NODE_IS_SAMPLE
+            )
+            if np.any(non_sample_nodes) and not include_non_sample_nodes:
+                individuals_nodes[non_sample_nodes] = -1
+                rows_to_reorder = np.any(non_sample_nodes, axis=1)
+                for i in np.where(rows_to_reorder)[0]:
+                    row = individuals_nodes[i]
+                    individuals_nodes[i] = np.concatenate(
+                        [row[row != -1], row[row == -1]]
                     )
-                ] = -1
 
             if individual_names is None:
                 if name_metadata_key is not None: