def test_bio_phylo_to_times(tree):
leaf_times, coal_times = bio_phylo_to_times(tree)
assert len(coal_times) + 1 == len(leaf_times)
# Check positivity.
times = torch.cat([coal_times, leaf_times])
signs = torch.cat([-torch.ones_like(coal_times), torch.ones_like(leaf_times)])
times, index = times.sort(0)
signs = signs[index]
lineages = signs.flip([0]).cumsum(0).flip([0])
assert (lineages >= 0).all()
def test_bio_phylo_to_times_custom(tree):
# Test a custom time parser.
def get_time(clade):
date_string = re.search(r"date=(\d\d\d\d\.\d\d)", clade.comment).group(1)
return (float(date_string) - 2020) * 365.25
leaf_times, coal_times = bio_phylo_to_times(tree, get_time=get_time)
assert len(coal_times) + 1 == len(leaf_times)
# Check positivity.
times = torch.cat([coal_times, leaf_times])
signs = torch.cat([-torch.ones_like(coal_times), torch.ones_like(leaf_times)])
times, index = times.sort(0)
signs = signs[index]
lineages = signs.flip([0]).cumsum(0).flip([0])
assert (lineages >= 0).all()