New semantics after recent discussion

Author: benjeffery

python/tskit/trees.py

@@ -10533,6 +10533,7 @@ def map_to_vcf_model(
         ploidy=None,
         name_metadata_key=None,
         individual_names=None,
+        include_non_sample_nodes=None,
     ):
         """
         Maps the sample nodes in this tree sequence to a representation suitable for
@@ -10571,6 +10572,8 @@ def map_to_vcf_model(
             ID is specified, if a specified individual has no nodes, or if the number of
             individuals doesn't match the number of names.
         """
+        if include_non_sample_nodes is None:
+            include_non_sample_nodes = False
 
         if name_metadata_key is not None and individual_names is not None:
             raise ValueError(
@@ -10596,63 +10599,31 @@ def map_to_vcf_model(
             if ploidy is None:
                 ploidy = 1
             individuals_nodes = self.sample_nodes_by_ploidy(ploidy)
-            ts_individual_names = [f"tsk_{i}" for i in range(len(individuals_nodes))]
+            if individual_names is None:
+                individual_names = [f"tsk_{i}" for i in range(len(individuals_nodes))]
         else:
-            individuals_nodes = []
-            ts_individual_names = []
             if individuals is None:
-                for ind in self.individuals():
-                    if len(ind.nodes) == 0:
-                        warnings.warn(
-                            f"Individual {ind.id} has no nodes associated with it.",
-                            stacklevel=1,
-                        )
-                        continue
-                    is_sample = np.array(
-                        [self.nodes_flags[u] & tskit.NODE_IS_SAMPLE for u in ind.nodes]
+                individuals = np.arange(self.num_individuals, dtype=np.int32)
+            if min(individuals) < 0 or max(individuals) >= self.num_individuals:
+                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
                     )
-                    if all(is_sample):
-                        individuals_nodes.append(ind.nodes)
-                        if name_metadata_key is not None:
-                            ts_individual_names.append(ind.metadata[name_metadata_key])
-                        else:
-                            ts_individual_names.append(f"tsk_{ind.id}")
-                    elif all(np.logical_not(is_sample)):
-                        continue
-                    else:
-                        warnings.warn(
-                            f"Individual {ind.id} has both sample and non-sample nodes "
-                            "associated with it.",
-                            stacklevel=1,
-                        )
-            else:
-                for i in individuals:
-                    if i < 0 or i >= self.num_individuals:
-                        raise ValueError(f"Invalid individual ID {i}")
-                    ind = self.individual(i)
-                    if len(ind.nodes) == 0:
-                        raise ValueError(
-                            f"Individual {i} has no nodes associated with it."
-                        )
-                    individuals_nodes.append(ind.nodes)
-                    if name_metadata_key is not None:
-                        ts_individual_names.append(ind.metadata[name_metadata_key])
-                    else:
-                        ts_individual_names.append(f"tsk_{ind.id}")
+                ] = -1
 
-            if len(individuals_nodes) > 0:
-                max_nodes = max(len(nodes) for nodes in individuals_nodes)
-                result = np.full(
-                    (len(individuals_nodes), max_nodes), -1, dtype=np.int32
-                )
-                for i, nodes in enumerate(individuals_nodes):
-                    result[i, : len(nodes)] = nodes
-                individuals_nodes = result
-            else:
-                individuals_nodes = np.empty((0, 0), dtype=np.int32)
+            if individual_names is None:
+                if name_metadata_key is not None:
+                    individual_names = [
+                        self.individual(i).metadata[name_metadata_key]
+                        for i in individuals
+                    ]
+                else:
+                    individual_names = [f"tsk_{i}" for i in individuals]
 
-        if individual_names is None:
-            individual_names = ts_individual_names
         individual_names = np.array(individual_names, dtype=object)
 
         if len(individuals_nodes) != len(individual_names):