Merge pull request #663 from jeromekelleher/support-one-sample

Allow a subset with one sample

Author: jeromekelleher

tests/test_formats.py

@@ -1,5 +1,5 @@
 #
-# Copyright (C) 2018-2020 University of Oxford
+# Copyright (C) 2018-2022 University of Oxford
 #
 # This file is part of tsinfer.
 #
@@ -569,7 +569,6 @@ def test_str(self):
 
     def test_eq(self):
         ts = tsutil.get_example_ts(5, random_seed=3)
-        print(ts.num_sites)
         input_file = formats.SampleData(sequence_length=ts.sequence_length)
         self.verify_data_round_trip(ts, input_file)
         assert input_file == input_file
@@ -663,9 +662,6 @@ def test_insufficient_samples(self):
         with pytest.raises(ValueError):
             sample_data.add_site(position=0, alleles=["0", "1"], genotypes=[])
         sample_data = formats.SampleData(sequence_length=10)
-        with pytest.raises(ValueError):
-            sample_data.add_site(position=0, alleles=["0", "1"], genotypes=[0])
-        sample_data = formats.SampleData(sequence_length=10)
         sample_data.add_individual(ploidy=3)
         with pytest.raises(ValueError):
             sample_data.add_site(position=0, alleles=["0", "1"], genotypes=[0])
@@ -1420,6 +1416,19 @@ def verify_subset_data(self, source, individuals, sites):
                 j += 1
         assert j == len(sites)
 
+    def test_one_sample(self):
+        ts = tsutil.get_example_ts(10)
+        sd1 = formats.SampleData.from_tree_sequence(ts)
+        G1 = ts.genotype_matrix()
+        # Because this is a haploid tree sequence we can use the
+        # individual and sample IDs interchangably.
+        cols = [3]
+        rows = np.arange(ts.num_sites)
+        subset = sd1.subset(individuals=cols, sites=rows)
+        G2 = np.array([v.genotypes for v in subset.variants()])
+        assert np.array_equal(G1[rows][:, cols], G2)
+        self.verify_subset_data(sd1, cols, rows)
+
     def test_simple_case(self):
         ts = tsutil.get_example_ts(10)
         sd1 = formats.SampleData.from_tree_sequence(ts)

tsinfer/formats.py

@@ -1372,8 +1372,8 @@ def subset(self, individuals=None, sites=None, **kwargs):
                     ploidy=len(individual.samples),
                 )
             sample_selection = np.array(sample_selection, dtype=int)
-            if len(sample_selection) < 2:
-                raise ValueError("Must have at least two samples")
+            if len(sample_selection) < 1:
+                raise ValueError("Must have at least one sample")
             for variant in self.variants():
                 if variant.site.id in sites:
                     subset.add_site(
@@ -1576,8 +1576,8 @@ def encode_metadata(metadata, schema):
         return self
 
     def _alloc_site_writer(self):
-        if self.num_samples < 2:
-            raise ValueError("Must have at least 2 samples")
+        if self.num_samples < 1:
+            raise ValueError("Must have at least 1 sample")
         self.sites_genotypes.resize(0, self.num_samples)
         arrays = {
             "position": self.sites_position,