Remove contig_id

Author: benjeffery

CHANGELOG.md

@@ -12,9 +12,8 @@
 - If a mismatch ratio is provided to the `infer` command, it only applies during the
   `match_samples` phase ({issue}`980`, {pr}`981`, {user}`hyanwong`)
 
-- Allow a contig to be selected by name (`contig_id`), and get the `sequence_length`
-  of the contig associated with the unmasked sites, if contig lengths are provided
-  ({pr}`964`, {user}`hyanwong`)
+- Get the `sequence_length` of the contig associated with the unmasked sites,
+  if contig lengths are provided ({pr}`964`, {user}`hyanwong`, {user}`benjeffery`)
 
 **Fixes**
 

docs/usage.md

@@ -109,8 +109,7 @@ It is also possible to *completely* exclude sites and samples, by specifing a bo
 a mask value of `True` will be completely omitted both from inference and the final tree sequence.
 This can be useful, for example, if you wish to select only a subset of the chromosome for
 inference, e.g. to reduce computational load. You can also use it to subset inference to a
-particular contig, if your dataset contains multiple contigs (although this can be more easily
-done using the `contig_id` parameter). Note that if a `site_mask` is provided,
+particular contig, if your dataset contains multiple contigs. Note that if a `site_mask` is provided,
 the ancestral states array should only specify alleles for the unmasked sites.
 
 Below, for instance, is an example of including only sites up to position six in the contig

tests/test_variantdata.py

@@ -145,7 +145,7 @@ def test_variantdata_accessors(tmp_path, in_mem):
     assert vd.format_name == "tsinfer-variant-data"
     assert vd.format_version == (0, 1)
     assert vd.finalised
-    assert vd.sequence_length == ts.sequence_length + 1337
+    assert vd.sequence_length == ts.sequence_length
     assert vd.num_sites == ts.num_sites
     assert vd.sites_metadata_schema == ts.tables.sites.metadata_schema.schema
     assert vd.sites_metadata == [site.metadata for site in ts.sites()]
@@ -317,6 +317,11 @@ def test_simulate_genotype_call_dataset_length(tmp_path):
     vdata = tsinfer.VariantData(tmp_path, ds["variant_allele"][:, 0].values.astype(str))
     assert vdata.sequence_length == ts.sites_position[-1] + 1
 
+    vdata = tsinfer.VariantData(
+        tmp_path, ds["variant_allele"][:, 0].values.astype(str), sequence_length=1337
+    )
+    assert vdata.sequence_length == 1337
+
 
 class TestMultiContig:
     def make_two_ts_dataset(self, path):
@@ -359,37 +364,44 @@ def test_mask(self, tmp_path):
         assert vdata.sequence_length == ts2.sequence_length
 
     @pytest.mark.parametrize("contig_id", ["chr1", "chr2"])
-    def test_contig_id_param(self, contig_id, tmp_path):
+    def test_multi_contig(self, contig_id, tmp_path):
         tree_seqs = {}
         tree_seqs["chr1"], tree_seqs["chr2"] = self.make_two_ts_dataset(tmp_path)
+        with pytest.raises(ValueError, match="multiple contigs"):
+            vdata = tsinfer.VariantData(tmp_path, "variant_ancestral_allele")
+        root = zarr.open(tmp_path)
+        mask = root["variant_contig"][:] == (1 if contig_id == "chr1" else 0)
         vdata = tsinfer.VariantData(
-            tmp_path, "variant_ancestral_allele", contig_id=contig_id
+            tmp_path, "variant_ancestral_allele", site_mask=mask
         )
         assert np.all(tree_seqs[contig_id].sites_position == vdata.sites_position)
         assert vdata.contig_id == contig_id
+        assert vdata._contig_index == (0 if contig_id == "chr1" else 1)
         assert vdata.sequence_length == tree_seqs[contig_id].sequence_length
 
-    def test_contig_id_param_and_mask(self, tmp_path):
+    def test_mixed_contigs_error(self, tmp_path):
         ts1, ts2 = self.make_two_ts_dataset(tmp_path)
-        vdata = tsinfer.VariantData(
-            tmp_path,
-            "variant_ancestral_allele",
-            site_mask=np.array(
-                (ts1.num_sites + 1) * [True] + (ts2.num_sites - 1) * [False]
-            ),
-            contig_id="chr2",
-        )
-        assert np.all(ts2.sites_position[1:] == vdata.sites_position)
-        assert vdata.contig_id == "chr2"
+        mask = np.ones(ts1.num_sites + ts2.num_sites)
+        # Select two varaints, one from each contig
+        mask[0] = False
+        mask[-1] = False
+        with pytest.raises(ValueError, match="multiple contigs"):
+            tsinfer.VariantData(
+                tmp_path,
+                "variant_ancestral_allele",
+                site_mask=mask,
+            )
 
-    @pytest.mark.parametrize("contig_id", ["chr1", "chr2"])
-    def test_contig_length(self, contig_id, tmp_path):
-        tree_seqs = {}
-        tree_seqs["chr1"], tree_seqs["chr2"] = self.make_two_ts_dataset(tmp_path)
+    def test_no_variant_contig(self, tmp_path):
+        ts1, ts2 = self.make_two_ts_dataset(tmp_path)
+        root = zarr.open(tmp_path)
+        del root["variant_contig"]
+        mask = np.ones(ts1.num_sites + ts2.num_sites)
+        mask[0] = False
         vdata = tsinfer.VariantData(
-            tmp_path, "variant_ancestral_allele", contig_id=contig_id
+            tmp_path, "variant_ancestral_allele", site_mask=mask
         )
-        assert vdata.sequence_length == tree_seqs[contig_id].sequence_length
+        assert vdata.sequence_length == ts1.sites_position[0] + 1
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="File permission errors on Windows")
@@ -953,23 +965,6 @@ def test_unimplemented_from_tree_sequence(self):
         with pytest.raises(NotImplementedError):
             tsinfer.VariantData.from_tree_sequence(None)
 
-    def test_multiple_contigs(self, tmp_path):
-        path = tmp_path / "data.zarr"
-        ds = sgkit.simulate_genotype_call_dataset(n_variant=3, n_sample=3, phased=True)
-        ds["contig_id"] = (
-            ds["contig_id"].dims,
-            np.array(["c10", "c11"], dtype="<U3"),
-        )
-        ds["variant_contig"] = (
-            ds["variant_contig"].dims,
-            np.array([0, 0, 1], dtype=ds["variant_contig"].dtype),
-        )
-        sgkit.save_dataset(ds, path)
-        with pytest.raises(
-            ValueError, match=r'Sites belong to multiple contigs \("c10", "c11"\)'
-        ):
-            tsinfer.VariantData(path, ds["variant_allele"][:, 0].astype(str))
-
     def test_all_masked(self, tmp_path):
         path = tmp_path / "data.zarr"
         ds = sgkit.simulate_genotype_call_dataset(n_variant=3, n_sample=3, phased=True)
@@ -979,28 +974,6 @@ def test_all_masked(self, tmp_path):
                 path, ds["variant_allele"][:, 0].astype(str), site_mask=np.ones(3, bool)
             )
 
-    def test_bad_contig_param(self, tmp_path):
-        path = tmp_path / "data.zarr"
-        ds = sgkit.simulate_genotype_call_dataset(n_variant=3, n_sample=3, phased=True)
-        sgkit.save_dataset(ds, path)
-        with pytest.raises(ValueError, match='"XX" not found'):
-            tsinfer.VariantData(
-                path, ds["variant_allele"][:, 0].astype(str), contig_id="XX"
-            )
-
-    def test_multiple_contig_param(self, tmp_path):
-        path = tmp_path / "data.zarr"
-        ds = sgkit.simulate_genotype_call_dataset(n_variant=3, n_sample=3, phased=True)
-        ds["contig_id"] = (
-            ds["contig_id"].dims,
-            np.array(["chr1", "chr1"], dtype="<U4"),
-        )
-        sgkit.save_dataset(ds, path)
-        with pytest.raises(ValueError, match='Multiple contigs named "chr1"'):
-            tsinfer.VariantData(
-                path, ds["variant_allele"][:, 0].astype(str), contig_id="chr1"
-            )
-
     def test_missing_sites_time(self, tmp_path):
         path = tmp_path / "data.zarr"
         ds = sgkit.simulate_genotype_call_dataset(n_variant=3, n_sample=3, phased=True)

tests/tsutil.py

@@ -339,11 +339,6 @@ def _make_ts_and_zarr(path, add_optional=False, shuffle_alleles=True):
         )
 
     if add_optional:
-        add_attribute_to_dataset(
-            "sequence_length",
-            ts.sequence_length + 1337,
-            path / "data.zarr",
-        )
         sites_md = tables.sites.metadata
         sites_md_offset = tables.sites.metadata_offset
         add_array_to_dataset(

tsinfer/formats.py

@@ -2328,8 +2328,7 @@ class VariantData(SampleData):
     :param Union(array, str) site_mask: A numpy array of booleans specifying which
         sites to mask out (exclude) from the dataset. Alternatively, a string
         can be provided, giving the name of an array in the input dataset which contains
-        the site mask. If ``None`` (default), all sites are included (unless restricted
-        to a ``contig_id``, see below).
+        the site mask. If ``None`` (default), all sites are included.
     :param Union(array, str) sites_time: A numpy array of floats specifying the relative
         time of occurrence of the mutation to the derived state at each site. This must
         be of the same length as the number of unmasked sites. Alternatively, a
@@ -2339,11 +2338,11 @@ class VariantData(SampleData):
         reasonable approximation to the relative order of ancestors used for inference.
         Time values are ignored for sites not used in inference, such as singletons,
         sites with more than two alleles, or sites with an unknown ancestral state.
-    :param str contig_id: The name of the contig to use (e.g. "chr1"), if the .vcz file
-        contains multiple contigs; contig names can be found in the `.contig_id array
-        of the input dataset. If provided, sites associated with any other contigs will
-        be added to the sites that are masked out. If ``None`` (default), do not mark
-        out sites on the basis of their contig ID.
+    :param int sequence_length: An integer specifying the resulting `sequence_length`
+        attribute of the output tree sequence. If not specified the `contig_length`
+        attribute from the undelying zarr store for the contig of the used variants.
+        If that is not present then the maximum position plus one of the used variants
+        is used.
     """
 
     FORMAT_NAME = "tsinfer-variant-data"
@@ -2357,8 +2356,10 @@ def __init__(
         sample_mask=None,
         site_mask=None,
         sites_time=None,
-        contig_id=None,
+        sequence_length=None,
     ):
+        self._sequence_length = sequence_length
+        self._contig_index = None
         try:
             if len(path_or_zarr.call_genotype.shape) == 3:
                 # Assumed to be a VCF Zarr hierarchy
@@ -2391,24 +2392,15 @@ def __init__(
             raise ValueError(
                 "Site mask array must be the same length as the number of unmasked sites"
             )
-        if contig_id is not None:
-            contig_index = np.where(self.data.contig_id[:] == contig_id)[0]
-            if len(contig_index) == 0:
-                raise ValueError(
-                    f'"{contig_id}" not found among the available contig IDs: '
-                    + ",".join(f"{n}" for n in self.data.contig_id[:])
-                )
-            elif len(contig_index) > 1:
-                raise ValueError(f'Multiple contigs named "{contig_id}"')
-            contig_index = contig_index[0]
-            site_mask = np.logical_or(
-                site_mask, self.data["variant_contig"][:] != contig_index
-            )
 
         # We negate the mask as it is much easier in numpy to have True=keep
         self.sites_select = ~site_mask.astype(bool)
+
         if np.sum(self.sites_select) == 0:
-            raise ValueError("All sites have been masked out. Please unmask some")
+            raise ValueError(
+                "All sites have been masked out, at least one value"
+                "must be 'False' in the site mask"
+            )
 
         if sample_mask is None:
             sample_mask = np.full(self._num_individuals_before_mask, False, dtype=bool)
@@ -2439,18 +2431,19 @@ def __init__(
                     " unphased"
                 )
 
-        used_contigs = self.data.variant_contig[:][self.sites_select]
-        self._contig_index = used_contigs[0]
-        self._contig_id = self.data.contig_id[self._contig_index]
+        if "variant_contig" in self.data:
+            used_contigs = self.data.variant_contig[:][self.sites_select]
+            self._contig_index = used_contigs[0]
+            self._contig_id = self.data.contig_id[self._contig_index]
 
-        if np.any(used_contigs != self._contig_index):
-            contig_names = ", ".join(
-                f'"{self.data.contig_id[c]}"' for c in np.unique(used_contigs)
-            )
-            raise ValueError(
-                f"Sites belong to multiple contigs ({contig_names}). Please restrict "
-                "sites to one contig e.g. via the `contig_id` argument."
-            )
+            if np.any(used_contigs != self._contig_index):
+                contig_names = ", ".join(
+                    f'"{self.data.contig_id[c]}"' for c in np.unique(used_contigs)
+                )
+                raise ValueError(
+                    f"Sites belong to multiple contigs ({contig_names}). "
+                    "Please restrict sites to one contig using the sites_mask argument."
+                )
 
         if np.any(np.diff(self.sites_position) <= 0):
             raise ValueError(
@@ -2558,26 +2551,21 @@ def finalised(self):
     def sequence_length(self):
         """
         The sequence length of the contig associated with sites used in the dataset.
-        If the dataset has a "sequence_length" attribute, this is always used, otherwise
-        if the dataset has recorded contig lengths, the appropriate length is taken,
-        otherwise the length is calculated from the maximum variant position plus one.
+        If set manually then that value is used else if the dataset has recorded
+        contig lengths use that else the length is calculated from the maximum
+        variant position plus one.
         """
-        try:
-            return self.data.attrs["sequence_length"]
-        except KeyError:
-            if self._contig_index is not None:
-                try:
-                    if self._contig_index < len(self.data.contig_length):
-                        return self.data.contig_length[self._contig_index]
-                except AttributeError:
-                    pass  # contig_length is optional, fall back to calculating length
-        return int(np.max(self.data["variant_position"])) + 1
+        if self._sequence_length is not None:
+            return self._sequence_length
+        if self._contig_index is not None and "contig_length" in self.data:
+            return self.data.contig_length[self._contig_index]
+        return int(np.max(self.sites_position)) + 1
 
     @property
     def contig_id(self):
         """
         The contig ID (name) for all used sites, or None if no
-        contig IDs were provided
+        contig IDs were present in the zarr dataset
         """
         return self._contig_id