Add fields to vcf model

Author: benjeffery

python/CHANGELOG.rst

@@ -13,6 +13,9 @@
 - ``draw_svg()`` methods now allow the y-axis to be placed on the right-hand side
   using ``y_axis="right"`` (:user:`hyanwong`, :pr:`3201`)
 
+- Add ``contig_id`` and ``isolated_as_missing`` to ``VcfModelMapping``
+  (:user:`benjeffery`, :pr:`3219`, :issue:`3177`)
+
 **Bugfixes**
 
 - Fix bug in ``TreeSequence.pair_coalescence_counts`` when ``span_normalise=True``

python/tests/test_highlevel.py

@@ -5969,3 +5969,27 @@ def bad_transform(positions):
             match="Position transform must return an array of the same length",
         ):
             ts.map_to_vcf_model(position_transform=bad_transform)
+
+    def test_contig_id(self):
+        tables = tskit.TableCollection(10.0)
+        tables.sites.add_row(position=1.0, ancestral_state="A")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        ts = tables.tree_sequence()
+
+        result = ts.map_to_vcf_model(contig_id="chr1")
+        assert result.contig_id == "chr1"
+
+        result = ts.map_to_vcf_model()
+        assert result.contig_id == "1"
+
+    def test_isolated_as_missing(self):
+        tables = tskit.TableCollection(10.0)
+        tables.sites.add_row(position=1.0, ancestral_state="A")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        ts = tables.tree_sequence()
+
+        result = ts.map_to_vcf_model(isolated_as_missing=False)
+        assert result.isolated_as_missing is False
+
+        result = ts.map_to_vcf_model()
+        assert result.isolated_as_missing is True

python/tskit/trees.py

@@ -65,6 +65,8 @@ class VcfModelMapping:
     individuals_name: np.ndarray
     transformed_positions: np.ndarray
     contig_length: int
+    contig_id: str
+    isolated_as_missing: bool
 
 
 class CoalescenceRecord(NamedTuple):
@@ -10645,6 +10647,8 @@ def map_to_vcf_model(
         individual_names=None,
         include_non_sample_nodes=None,
         position_transform=None,
+        contig_id=None,
+        isolated_as_missing=None,
     ):
         """
         Maps the sample nodes in this tree sequence to a representation suitable for
@@ -10700,6 +10704,12 @@ def map_to_vcf_model(
             by incrementing is used.
             See the :ref:`sec_export_vcf_modifying_coordinates` for examples
             and more information.
+        :param str contig_id: The ID of the contig to use in the VCF output.
+            Defaults to "1" if not specified.
+        :param bool isolated_as_missing: If True, isolated samples without mutations
+            will be considered as missing data in the VCF output. If False, these samples
+            will have the ancestral state in the VCF output.
+            Default: True.
         :return: A VcfModelMapping containing the node-to-individual mapping,
             individual names, transformed positions, and transformed contig length.
         :raises ValueError: If both name_metadata_key and individual_names are specified,
@@ -10710,6 +10720,12 @@ def map_to_vcf_model(
         if include_non_sample_nodes is None:
             include_non_sample_nodes = False
 
+        if contig_id is None:
+            contig_id = "1"
+
+        if isolated_as_missing is None:
+            isolated_as_missing = True
+
         if name_metadata_key is not None and individual_names is not None:
             raise ValueError(
                 "Cannot specify both name_metadata_key and individual_names"
@@ -10808,7 +10824,12 @@ def legacy_position_transform(positions):
         contig_length = max(1, int(position_transform([self.sequence_length])[0]))
 
         return VcfModelMapping(
-            individuals_nodes, individual_names, transformed_positions, contig_length
+            individuals_nodes,
+            individual_names,
+            transformed_positions,
+            contig_length,
+            contig_id,
+            isolated_as_missing,
         )
 
     ############################################

python/tskit/vcf.py

@@ -50,62 +50,59 @@ def __init__(
         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,
             include_non_sample_nodes=include_non_sample_nodes,
             position_transform=position_transform,
+            contig_id=contig_id,
+            isolated_as_missing=isolated_as_missing,
         )
+
+        # We now make some tweaks to the VCF model required for
+        # writing the VCF in text format
+
         # 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]
+        to_keep = (vcf_model.individuals_nodes != -1).any(axis=1)
+        vcf_model.individuals_nodes = vcf_model.individuals_nodes[to_keep]
+        vcf_model.individual_names = vcf_model.individuals_name[to_keep]
         self.individual_ploidies = [
-            len(nodes[nodes >= 0]) for nodes in individuals_nodes
+            len(nodes[nodes >= 0]) for nodes in vcf_model.individuals_nodes
         ]
-        self.num_individuals = len(self.individual_names)
+        self.num_individuals = len(vcf_model.individual_names)
 
-        if len(individuals_nodes) == 0:
+        if len(vcf_model.individuals_nodes) == 0:
             raise ValueError("No samples in resulting VCF model")
 
+        if len(vcf_model.transformed_positions) > 0:
+            # Arguably this should be last_pos + 1, but if we hit this
+            # condition the coordinate systems are all muddled up anyway
+            # so it's simpler to stay with this rule that was inherited
+            # from the legacy VCF output code.
+            vcf_model.contig_length = max(
+                vcf_model.transformed_positions[-1], vcf_model.contig_length
+            )
+
         # Flatten the array of node IDs, filtering out the -1 padding values
         self.samples = []
-        for row in individuals_nodes:
+        for row in vcf_model.individuals_nodes:
             for node_id in row:
                 if node_id != -1:
                     self.samples.append(node_id)
 
-        self.transformed_positions = vcf_model.transformed_positions
-        self.contig_length = vcf_model.contig_length
-        if len(self.transformed_positions) > 0:
-            # Arguably this should be last_pos + 1, but if we hit this
-            # condition the coordinate systems are all muddled up anyway
-            # so it's simpler to stay with this rule that was inherited
-            # from the legacy VCF output code.
-            self.contig_length = max(self.transformed_positions[-1], self.contig_length)
-
         if site_mask is None:
             site_mask = np.zeros(tree_sequence.num_sites, dtype=bool)
         self.site_mask = np.array(site_mask, dtype=bool)
         if self.site_mask.shape != (tree_sequence.num_sites,):
             raise ValueError("Site mask must be 1D a boolean array of length num_sites")
 
-        self.sample_mask = sample_mask
-        if sample_mask is not None:
-            if not callable(sample_mask):
-                sample_mask = np.array(sample_mask, dtype=bool)
-                self.sample_mask = lambda _: sample_mask
-
         # The VCF spec does not allow for positions to be 0, so we error if one of the
         # transformed positions is 0 and allow_position_zero is False.
         if not allow_position_zero and np.any(
-            self.transformed_positions[~site_mask] == 0
+            vcf_model.transformed_positions[~site_mask] == 0
         ):
             raise ValueError(
                 "A variant position of 0 was found in the VCF output, this is not "
@@ -116,17 +113,26 @@ def __init__(
                 '"position_transform = lambda x: np.fmax(1, x)"'
             )
 
+        self.sample_mask = sample_mask
+        if sample_mask is not None:
+            if not callable(sample_mask):
+                sample_mask = np.array(sample_mask, dtype=bool)
+                self.sample_mask = lambda _: sample_mask
+
+        self.vcf_model = vcf_model
+
     def __write_header(self, output):
         print("##fileformat=VCFv4.2", file=output)
         print(f"##source=tskit {provenance.__version__}", file=output)
         print('##FILTER=<ID=PASS,Description="All filters passed">', file=output)
         print(
-            f"##contig=<ID={self.contig_id},length={self.contig_length}>", file=output
+            f"##contig=<ID={self.vcf_model.contig_id},length={self.vcf_model.contig_length}>",
+            file=output,
         )
         print(
             '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">', file=output
         )
-        vcf_samples = "\t".join(self.individual_names)
+        vcf_samples = "\t".join(self.vcf_model.individual_names)
         print(
             "#CHROM",
             "POS",
@@ -163,7 +169,7 @@ def write(self, output):
         indexes = np.array(indexes, dtype=int)
 
         for variant in self.tree_sequence.variants(
-            samples=self.samples, isolated_as_missing=self.isolated_as_missing
+            samples=self.samples, isolated_as_missing=self.vcf_model.isolated_as_missing
         ):
             site_id = variant.site.id
             # We check the mask before we do any checks so we can use this as a
@@ -176,13 +182,13 @@ def write(self, output):
                     "More than 9 alleles not currently supported. Please open an issue "
                     "on GitHub if this limitation affects you."
                 )
-            pos = self.transformed_positions[variant.index]
+            pos = self.vcf_model.transformed_positions[variant.index]
             ref = variant.alleles[0]
             alt = "."
             if variant.num_alleles > 1:
                 alt = ",".join(variant.alleles[1 : variant.num_alleles])
             print(
-                self.contig_id,
+                self.vcf_model.contig_id,
                 pos,
                 site_id,
                 ref,