Allow haplotypes and alignments to have start, stop, samples params

Author: hyanwong

python/CHANGELOG.rst

@@ -14,6 +14,10 @@
 
 - New ``Site.alleles()`` method (:user:`hyanwong`, :issue:`2380`, :pr:`2385`)
 
+- The ``variants()``, ``haplotypes()`` and ``alignments()`` methods can now
+  take a list of sample ids and a left and right position, to restrict the
+  size of the output (:user:`hyanwong`, :issue:`2092`, :pr:`2397`)
+
 
 --------------------
 [0.5.0] - 2022-06-22

python/tests/test_genotypes.py

@@ -638,6 +638,33 @@ def test_snipped_tree_sequence_mutations_over_isolated(self):
         assert non_missing_found
         assert missing_found
 
+    def test_limit_interval(self):
+        ts = self.get_tree_sequence()
+        test_variant = tskit.Variant(ts)
+        test_variant.decode(1)
+        for v in ts.variants(left=ts.site(1).position, right=ts.site(2).position):
+            # should only decode the first variant
+            assert v.site.id == 1
+            assert np.all(v.genotypes == test_variant.genotypes)
+            assert v.alleles == test_variant.alleles
+
+    def test_bad_left(self):
+        ts = tskit.TableCollection(10).tree_sequence()
+        for bad_left in [-1, 10, 100, np.nan, np.inf, -np.inf]:
+            with pytest.raises(ValueError, match="`left` not between"):
+                list(ts.variants(left=bad_left))
+
+    def test_bad_right(self):
+        ts = tskit.TableCollection(10).tree_sequence()
+        for bad_right in [-1, 0, 100, np.nan, np.inf, -np.inf]:
+            with pytest.raises(ValueError, match="`right` not between"):
+                list(ts.variants(right=bad_right))
+
+    def test_bad_left_right(self):
+        ts = tskit.TableCollection(10).tree_sequence()
+        with pytest.raises(ValueError, match="must be less than"):
+            list(ts.variants(left=1, right=1))
+
 
 class TestHaplotypeGenerator:
     """
@@ -830,6 +857,29 @@ def test_missing_data(self):
             )
         assert h == ["A", "A"]
 
+    def test_restrict_samples(self):
+        tables = tskit.TableCollection(1.0)
+        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
+        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
+        tables.sites.add_row(0.5, "A")
+        tables.mutations.add_row(0, 0, derived_state="B")
+        ts = tables.tree_sequence()
+        haplotypes = list(ts.haplotypes(samples=[0], isolated_as_missing=False))
+        assert haplotypes == ["B"]
+        haplotypes = list(ts.haplotypes(samples=[1], isolated_as_missing=False))
+        assert haplotypes == ["A"]
+
+    def test_restrict_positions(self):
+        tables = tskit.TableCollection(1.0)
+        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
+        tables.sites.add_row(0.1, "A")
+        tables.sites.add_row(0.2, "B")
+        tables.sites.add_row(0.3, "C")
+        tables.sites.add_row(0.4, "D")
+        ts = tables.tree_sequence()
+        haplotypes = list(ts.haplotypes(left=0.2, right=0.4, isolated_as_missing=False))
+        assert haplotypes == ["BC"]
+
 
 class TestUserAlleles:
     """
@@ -1091,6 +1141,14 @@ def test_alignments_default(self):
         assert A[1] == "NNANNNNNNC"
         assert A[2] == "NNANNNNNNC"
 
+    def test_alignments_restricted(self):
+        ts = self.ts()
+        samples = ts.samples()
+        # Take the first 2 in reverse order
+        A = list(ts.alignments(left=1, right=9, samples=samples[1::-1]))
+        assert A[0] == "NANNNNNN"
+        assert A[1] == "NGNNNNNN"
+
     def test_alignments_missing_data_char(self):
         A = list(self.ts().alignments(missing_data_character="x"))
         assert A[0] == "xxGxxxxxxT"
@@ -1669,13 +1727,17 @@ def test_reference_length_mismatch(self, ref_length):
         tables = tskit.TableCollection(10)
         tables.reference_sequence.data = "A" * ref_length
         ts = tables.tree_sequence()
-        with pytest.raises(ValueError, match="same length"):
+        if ref_length <= tables.sequence_length:
+            with pytest.raises(ValueError, match="shorter than"):
+                list(ts.alignments())
+        else:
+            # Longer reference sequences are allowed
             list(ts.alignments())
 
     @pytest.mark.parametrize("ref", ["", "xy"])
     def test_reference_sequence_length_mismatch(self, ref):
         ts = self.simplest_ts()
-        with pytest.raises(ValueError, match="same length"):
+        with pytest.raises(ValueError, match="shorter than"):
             list(ts.alignments(reference_sequence=ref))
 
     @pytest.mark.parametrize("ref", ["À", "┃", "α"])
@@ -1699,6 +1761,23 @@ def test_non_ascii_missing_data_char(self, missing_data_char):
         with pytest.raises(UnicodeEncodeError):
             list(ts.alignments(missing_data_character=missing_data_char))
 
+    def test_bad_left(self):
+        ts = tskit.TableCollection(10).tree_sequence()
+        with pytest.raises(ValueError, match="integer"):
+            list(ts.alignments(left=0.1))
+
+    def test_bad_right(self):
+        ts = tskit.TableCollection(10).tree_sequence()
+        with pytest.raises(ValueError, match="integer"):
+            list(ts.alignments(right=1.1))
+
+    def test_bad_restricted(self):
+        tables = tskit.TableCollection(10)
+        tables.reference_sequence.data = "A" * 7
+        ts = tables.tree_sequence()
+        with pytest.raises(ValueError, match="sequence ends before"):
+            list(ts.alignments(right=8))
+
 
 class TestAlignmentExamples:
     @pytest.mark.parametrize("ts", get_example_discrete_genome_tree_sequences())