Add sample mask

Author: benjeffery

tests/test_sgkit.py

@@ -481,9 +481,50 @@ def test_bad_mask_length_at_iterator(self, tmp_path):
         with pytest.raises(
             ValueError, match="Mask must be the same length as the array"
         ):
-            for _ in chunk_iterator(ds.variant_position, mask=sites_mask):
+            for _ in chunk_iterator(ds.call_genotype, mask=sites_mask):
                 pass
 
+    @pytest.mark.parametrize("sample_list", [[1, 2, 3, 5, 9, 27], [0], []])
+    def test_sgkit_sample_mask(self, tmp_path, sample_list):
+        ts, zarr_path = make_ts_and_zarr(tmp_path, add_optional=True)
+        ds = sgkit.load_dataset(zarr_path)
+        samples_mask = np.zeros_like(ds["sample_id"], dtype=bool)
+        for i in sample_list:
+            samples_mask[i] = True
+        add_array_to_dataset("samples_mask", samples_mask, zarr_path)
+        samples = tsinfer.SgkitSampleData(zarr_path)
+        assert samples.ploidy == 3
+        assert samples.num_individuals == len(sample_list)
+        assert samples.num_samples == len(sample_list) * samples.ploidy
+        assert np.array_equal(samples.individuals_mask, samples_mask)
+        assert np.array_equal(samples.samples_mask, np.repeat(samples_mask, 3))
+        assert np.array_equal(
+            samples.individuals_time, ds.individuals_time.values[samples_mask]
+        )
+        assert np.array_equal(
+            samples.individuals_location, ds.individuals_location.values[samples_mask]
+        )
+        assert np.array_equal(
+            samples.individuals_population,
+            ds.individuals_population.values[samples_mask],
+        )
+        assert np.array_equal(
+            samples.individuals_flags, ds.individuals_flags.values[samples_mask]
+        )
+        assert np.array_equal(
+            samples.samples_individual, np.repeat(np.arange(len(sample_list)), 3)
+        )
+        expected_gt = ds.call_genotype.values[:, samples_mask, :].reshape(
+            samples.num_sites, len(sample_list) * 3
+        )
+        assert np.array_equal(samples.sites_genotypes, expected_gt)
+        for v, gt in zip(samples.variants(), expected_gt):
+            assert np.array_equal(v.genotypes, gt)
+
+        for i, (id, haplo) in enumerate(samples.haplotypes()):
+            assert id == i
+            assert np.array_equal(haplo, expected_gt[:, i])
+
 
 @pytest.mark.skipif(sys.platform == "win32", reason="No cyvcf2 on Windows")
 def test_sgkit_ancestral_allele_same_ancestors(tmp_path):

tsinfer/formats.py

@@ -305,7 +305,7 @@ def zarr_summary(array):
     return ret
 
 
-def chunk_iterator(array, indexes=None, mask=None, dimension=0):
+def chunk_iterator(array, indexes=None, mask=None, orthogonal_mask=None, dimension=0):
     """
     Utility to iterate over closely spaced rows in the specified array efficiently
     by accessing one chunk at a time (normally used as an iterator over each row)
@@ -314,6 +314,8 @@ def chunk_iterator(array, indexes=None, mask=None, dimension=0):
     assert dimension < 2
     if mask is None:
         mask = np.ones(array.shape[dimension], dtype=bool)
+    if orthogonal_mask is None:
+        orthogonal_mask = np.ones(array.shape[int(not dimension)], dtype=bool)
     if len(mask) != array.shape[dimension]:
         raise ValueError("Mask must be the same length as the array")
 
@@ -339,14 +341,14 @@ def chunk_iterator(array, indexes=None, mask=None, dimension=0):
             if chunk_id != prev_chunk_id:
                 chunk = array[chunk_id * chunk_size : (chunk_id + 1) * chunk_size][:]
                 prev_chunk_id = chunk_id
-            yield chunk[j % chunk_size]
+            yield chunk[j % chunk_size, orthogonal_mask]
     elif dimension == 1:
         for j in indexes:
             chunk_id = j // chunk_size
             if chunk_id != prev_chunk_id:
                 chunk = array[:, chunk_id * chunk_size : (chunk_id + 1) * chunk_size][:]
                 prev_chunk_id = chunk_id
-            yield chunk[:, j % chunk_size]
+            yield chunk[orthogonal_mask, j % chunk_size]
 
 
 def merge_variants(sd1, sd2):
@@ -2297,9 +2299,9 @@ def __init__(self, path):
         self.path = path
         self.data = zarr.open(path, mode="r")
         genotypes_arr = self.data["call_genotype"]
-        _, self._num_individuals, self.ploidy = genotypes_arr.shape
+        _, self._num_unmasked_individuals, self.ploidy = genotypes_arr.shape
         self._num_sites = np.sum(self.sites_mask)
-        self._num_samples = self._num_individuals * self.ploidy
+        self._num_unmasked_samples = self._num_unmasked_individuals * self.ploidy
 
         assert self.ploidy == self.data["call_genotype"].chunks[2]
         if self.ploidy > 1:
@@ -2333,6 +2335,19 @@ def sequence_length(self):
     def num_sites(self):
         return self._num_sites
 
+    @functools.cached_property
+    def individuals_mask(self):
+        try:
+            return self.data["samples_mask"][:].astype(bool)
+        except KeyError:
+            return np.full(self._num_unmasked_individuals, True, dtype=bool)
+
+    @functools.cached_property
+    def samples_mask(self):
+        # Samples in sgkit are individuals in tskit, so we need to expand
+        # the mask to cover all the samples for each individual.
+        return np.repeat(self.individuals_mask, self.ploidy)
+
     @functools.cached_property
     def sites_metadata_schema(self):
         try:
@@ -2427,9 +2442,9 @@ def sites_genotypes(self):
         gt = self.data["call_genotype"]
         # This method is only used for test/debug so we retrieve and
         # reshape the entire array.
-        return gt[...][self.sites_mask, :, :].reshape(
-            gt.shape[0], gt.shape[1] * gt.shape[2]
-        )
+        ret = gt[...][self.sites_mask, :, :]
+        ret = ret[:, self.individuals_mask, :]
+        return ret.reshape(ret.shape[0], ret.shape[1] * ret.shape[2])
 
     @functools.cached_property
     def provenances_timestamp(self):
@@ -2445,9 +2460,9 @@ def provenances_record(self):
         except KeyError:
             return np.array([], dtype=object)
 
-    @property
+    @functools.cached_property
     def num_samples(self):
-        return self._num_samples
+        return np.sum(self.samples_mask)
 
     @functools.cached_property
     def samples_individual(self):
@@ -2500,12 +2515,12 @@ def populations_metadata_schema(self):
 
     @property
     def num_individuals(self):
-        return self._num_individuals
+        return np.sum(self.individuals_mask)
 
     @functools.cached_property
     def individuals_time(self):
         try:
-            return self.data["individuals_time"]
+            return self.data["individuals_time"][:][self.individuals_mask]
         except KeyError:
             return np.full(self.num_individuals, tskit.UNKNOWN_TIME)
 
@@ -2524,11 +2539,14 @@ def individuals_metadata(self):
         # We set the sample_id in the individual metadata as this is often useful,
         # however we silently don't overwrite if the key exists
         if "individuals_metadata" in self.data:
-            assert len(self.data["individuals_metadata"]) == self.num_individuals
-            assert self.num_individuals == len(self.data["sample_id"])
+            assert (
+                len(self.data["individuals_metadata"]) == self._num_unmasked_individuals
+            )
+            assert self._num_unmasked_individuals == len(self.data["sample_id"])
             md_list = []
             for sample_id, r in zip(
-                self.data["sample_id"], self.data["individuals_metadata"][:]
+                self.data["sample_id"][:][self.individuals_mask],
+                self.data["individuals_metadata"][:][self.individuals_mask],
             ):
                 md = schema.decode_row(r)
                 if "sgkit_sample_id" not in md:
@@ -2537,27 +2555,28 @@ def individuals_metadata(self):
             return md_list
         else:
             return [
-                {"sgkit_sample_id": sample_id} for sample_id in self.data["sample_id"]
+                {"sgkit_sample_id": sample_id}
+                for sample_id in self.data["sample_id"][:][self.individuals_mask]
             ]
 
     @functools.cached_property
     def individuals_location(self):
         try:
-            return self.data["individuals_location"]
+            return self.data["individuals_location"][:][self.individuals_mask]
         except KeyError:
             return np.array([[]] * self.num_individuals, dtype=float)
 
     @functools.cached_property
     def individuals_population(self):
         try:
-            return self.data["individuals_population"]
+            return self.data["individuals_population"][:][self.individuals_mask]
         except KeyError:
             return np.full((self.num_individuals), tskit.NULL, dtype=np.int32)
 
     @functools.cached_property
     def individuals_flags(self):
         try:
-            return self.data["individuals_flags"]
+            return self.data["individuals_flags"][:][self.individuals_mask]
         except KeyError:
             return np.full((self.num_individuals), 0, dtype=np.int32)
 
@@ -2585,7 +2604,10 @@ def variants(self, sites=None, recode_ancestral=None):
         if recode_ancestral is None:
             recode_ancestral = False
         all_genotypes = chunk_iterator(
-            self.data["call_genotype"], indexes=sites, mask=self.sites_mask
+            self.data["call_genotype"],
+            indexes=sites,
+            mask=self.sites_mask,
+            orthogonal_mask=self.individuals_mask,
         )
         assert MISSING_DATA < 0  # required for geno_map to remap MISSING_DATA
         for genos, site in zip(all_genotypes, self.sites(ids=sites)):
@@ -2627,9 +2649,11 @@ def _all_haplotypes(self, sites=None, recode_ancestral=None):
         aa_index[aa_index == MISSING_DATA] = 0
         gt = self.data["call_genotype"]
         chunk_size = gt.chunks[1]
-        for j in range(self.num_individuals):
-            if j % chunk_size == 0:
-                chunk = gt[:, j : j + chunk_size, :]
+        current_chunk = None
+        for j in np.where(self.individuals_mask)[0]:
+            if j // chunk_size != current_chunk:
+                current_chunk = j // chunk_size
+                chunk = gt[:, j // chunk_size : (j // chunk_size) + chunk_size, :]
                 # Zarr doesn't support fancy indexing, so we have to do this after
                 chunk = chunk[self.sites_mask]
             indiv_gt = chunk[:, j % chunk_size, :]

tsinfer/inference.py

@@ -1644,6 +1644,7 @@ def group_by_linesweep(self):
         epoch_end = np.hstack([breaks + 1, [self.num_ancestors]])
         time_slices = np.vstack([epoch_start, epoch_end]).T
         epoch_sizes = time_slices[:, 1] - time_slices[:, 0]
+
         median_size = np.median(epoch_sizes)
         cutoff = 500 * median_size
         # Zero out the first half so that an initial large epoch doesn't