Allow sequence length to be specified in sd.subset

Author: hyanwong

CHANGELOG.rst

@@ -2,6 +2,10 @@
 [0.2.4] - 2022-06-xx
 ********************
 
+**Fixes**
+
+- sample_data.subset() now accepts a sequence_length  (:pr:`681`, :user:`hyanwong`)
+
 **Breaking changes**:
 
 - Inference now sets time_units on both ancestor and final tree sequences to

tests/test_formats.py

@@ -1515,6 +1515,19 @@ def test_file_kwargs(self):
             sd2 = formats.SampleData.load(path)
             assert sd1.data_equal(sd2)
 
+    def test_sequence_length_change(self):
+        ts = tsutil.get_example_ts(10)
+        sd1 = formats.SampleData.from_tree_sequence(ts)
+        max_site_to_use = sd1.num_sites // 2
+        new_seq_length = sd1.sites_position[max_site_to_use + 1]
+        assert new_seq_length < sd1.sequence_length
+        assert max_site_to_use > 0
+        sd2 = sd1.subset(
+            sites=np.arange(max_site_to_use),
+            sequence_length=new_seq_length,
+        )
+        assert sd2.sequence_length == new_seq_length
+
 
 class TestSampleDataMerge:
     """

tsinfer/formats.py

@@ -1314,7 +1314,7 @@ def assert_data_equal(self, other):
         assert self.samples_equal(other)
         assert self.sites_equal(other)
 
-    def subset(self, individuals=None, sites=None, **kwargs):
+    def subset(self, individuals=None, sites=None, *, sequence_length=None, **kwargs):
         """
         Returns a subset of this sample data file consisting of the specified
         individuals and sites. It is important to note that these are
@@ -1356,7 +1356,9 @@ def subset(self, individuals=None, sites=None, **kwargs):
         sites = set(sites)
         if len(sites) != num_sites:
             raise ValueError("Duplicate site IDS")
-        with SampleData(sequence_length=self.sequence_length, **kwargs) as subset:
+        if sequence_length is None:
+            sequence_length = self.sequence_length
+        with SampleData(sequence_length=sequence_length, **kwargs) as subset:
             # NOTE We don't bother filtering the populations, but we could.
             for population in self.populations():
                 subset.add_population(population.metadata)