Add mask names

Author: benjeffery

tests/test_sgkit.py

@@ -438,8 +438,8 @@ def test_sgkit_variant_mask(self, tmp_path, sites):
         sites_mask = np.ones_like(ds["variant_position"], dtype=bool)
         for i in sites:
             sites_mask[i] = False
-        add_array_to_dataset("variant_mask", sites_mask, zarr_path)
-        samples = tsinfer.SgkitSampleData(zarr_path)
+        add_array_to_dataset("variant_mask_42", sites_mask, zarr_path)
+        samples = tsinfer.SgkitSampleData(zarr_path, sites_mask_name="variant_mask_42")
         assert samples.num_sites == len(sites)
         assert np.array_equal(samples.sites_mask, ~sites_mask)
         assert np.array_equal(
@@ -465,12 +465,13 @@ def test_sgkit_variant_bad_mask_length(self, tmp_path):
         ts, zarr_path = make_ts_and_zarr(tmp_path)
         ds = sgkit.load_dataset(zarr_path)
         sites_mask = np.zeros(ds.sizes["variants"] + 1, dtype=int)
-        add_array_to_dataset("variant_mask", sites_mask, zarr_path)
+        add_array_to_dataset("variant_mask_foobar", sites_mask, zarr_path)
+        tsinfer.SgkitSampleData(zarr_path)
         with pytest.raises(
             ValueError,
             match="Mask must be the same length as the number of unmasked sites",
         ):
-            tsinfer.SgkitSampleData(zarr_path)
+            tsinfer.SgkitSampleData(zarr_path, sites_mask_name="variant_mask_foobar")
 
     def test_bad_mask_length_at_iterator(self, tmp_path):
         ts, zarr_path = make_ts_and_zarr(tmp_path)
@@ -491,8 +492,10 @@ def test_sgkit_sample_mask(self, tmp_path, sample_list):
         samples_mask = np.ones_like(ds["sample_id"], dtype=bool)
         for i in sample_list:
             samples_mask[i] = False
-        add_array_to_dataset("samples_mask", samples_mask, zarr_path)
-        samples = tsinfer.SgkitSampleData(zarr_path)
+        add_array_to_dataset("samples_mask_69", samples_mask, zarr_path)
+        samples = tsinfer.SgkitSampleData(
+            zarr_path, sgkit_samples_mask_name="samples_mask_69"
+        )
         assert samples.ploidy == 3
         assert samples.num_individuals == len(sample_list)
         assert samples.num_samples == len(sample_list) * samples.ploidy
@@ -525,6 +528,24 @@ def test_sgkit_sample_mask(self, tmp_path, sample_list):
             assert id == i
             assert np.array_equal(haplo, expected_gt[:, i])
 
+    def test_sgkit_missing_masks(self, tmp_path):
+        ts, zarr_path = make_ts_and_zarr(tmp_path)
+        samples = tsinfer.SgkitSampleData(zarr_path)
+        samples.individuals_mask
+        samples.sites_mask
+        with pytest.raises(
+            ValueError, match="The sites mask foobar was not found in the dataset."
+        ):
+            tsinfer.SgkitSampleData(zarr_path, sites_mask_name="foobar")
+        with pytest.raises(
+            ValueError,
+            match="The sgkit samples mask foobar2 was not found in the dataset.",
+        ):
+            samples = tsinfer.SgkitSampleData(
+                zarr_path, sgkit_samples_mask_name="foobar2"
+            )
+            samples.individuals_mask
+
 
 @pytest.mark.skipif(sys.platform == "win32", reason="No cyvcf2 on Windows")
 def test_sgkit_ancestral_allele_same_ancestors(tmp_path):

tsinfer/formats.py

@@ -2295,9 +2295,11 @@ class SgkitSampleData(SampleData):
     FORMAT_NAME = "tsinfer-sgkit-sample-data"
     FORMAT_VERSION = (0, 1)
 
-    def __init__(self, path):
+    def __init__(self, path, sgkit_samples_mask_name=None, sites_mask_name=None):
         self.path = path
         self.data = zarr.open(path, mode="r")
+        self._sgkit_samples_mask_name = sgkit_samples_mask_name
+        self._sites_mask_name = sites_mask_name
         genotypes_arr = self.data["call_genotype"]
         _, self._num_unmasked_individuals, self.ploidy = genotypes_arr.shape
         self._num_sites = np.sum(self.sites_mask)
@@ -2337,11 +2339,17 @@ def num_sites(self):
 
     @functools.cached_property
     def individuals_mask(self):
-        try:
-            # We negate the mask as it is much easier in numpy to have True=keep
-            return ~(self.data["samples_mask"][:].astype(bool))
-        except KeyError:
+        if self._sgkit_samples_mask_name is None:
             return np.full(self._num_unmasked_individuals, True, dtype=bool)
+        else:
+            try:
+                # We negate the mask as it is much easier in numpy to have True=keep
+                return ~(self.data[self._sgkit_samples_mask_name][:].astype(bool))
+            except KeyError:
+                raise ValueError(
+                    f"The sgkit samples mask {self._sgkit_samples_mask_name} was not"
+                    f" found in the dataset."
+                )
 
     @functools.cached_property
     def samples_mask(self):
@@ -2386,18 +2394,24 @@ def sites_alleles(self):
 
     @functools.cached_property
     def sites_mask(self):
-        try:
-            if (
-                self.data["variant_mask"].shape[0]
-                != self.data["variant_position"].shape[0]
-            ):
+        if self._sites_mask_name is None:
+            return np.full(self.data["variant_position"].shape, True, dtype=bool)
+        else:
+            try:
+                if (
+                    self.data[self._sites_mask_name].shape[0]
+                    != self.data["variant_position"].shape[0]
+                ):
+                    raise ValueError(
+                        "Mask must be the same length as the number of unmasked sites"
+                    )
+                # We negate the mask as it is much easier in numpy to have True=keep
+                return ~(self.data[self._sites_mask_name].astype(bool)[:])
+            except KeyError:
                 raise ValueError(
-                    "Mask must be the same length as the number of unmasked sites"
+                    f"The sites mask {self._sites_mask_name} was not found"
+                    f" in the dataset."
                 )
-            # We negate the mask as it is much easier in numpy to have True=keep
-            return ~(self.data["variant_mask"].astype(bool)[:])
-        except KeyError:
-            return np.full(self.data["variant_position"].shape, True, dtype=bool)
 
     @functools.cached_property
     def sites_ancestral_allele(self):