Skip to content

Commit

Permalink
Allow sequence length to be specified in sd.subset
Browse files Browse the repository at this point in the history
  • Loading branch information
hyanwong committed Jun 21, 2022
1 parent 4f76b6a commit 75efd0b
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 1 deletion.
4 changes: 4 additions & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
13 changes: 13 additions & 0 deletions tests/test_formats.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down
4 changes: 3 additions & 1 deletion tsinfer/formats.py
Original file line number Diff line number Diff line change
Expand Up @@ -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" not in kwargs:
kwargs["sequence_length"] = self.sequence_length
with SampleData(**kwargs) as subset:
# NOTE We don't bother filtering the populations, but we could.
for population in self.populations():
subset.add_population(population.metadata)
Expand Down

0 comments on commit 75efd0b

Please sign in to comment.