Implementation:Pyro ppl Pyro CoalescentTimes

Knowledge Sources	Pyro Pyro Distributions
Domains	Probability_Distributions, Bayesian_Phylogenetics
Last Updated	2026-02-09 09:00 GMT

Overview

Distributions and utilities for modeling coalescent times in phylogenetic trees under the Kingman coalescent process.

Description

This module implements several classes for coalescent theory in population genetics and phylogenetics. The core distribution, CoalescentTimes, models sorted binary coalescent event times given irregular leaf sampling times and a constant population size. It extends TorchDistribution and produces samples that are sorted sets of coalescent times with cardinality equal to the number of leaves minus one, forming complete binary trees.

The module provides three main components:

• CoalescentTimes -- A full distribution (with both sample and log_prob) for coalescent times under a constant-rate Kingman coalescent model. The log probability is computed by treating the coalescent as a Poisson process with rate binomial in the number of active lineages.

• CoalescentTimesWithRate -- A distribution for coalescent times given piecewise constant coalescent rates on a regular time grid. This supports only log_prob (not sampling) and computes the likelihood following Popinga et al. (2014), equations 7-8. The time complexity is O(T + S N log(N)) where T is the number of time steps, N is the number of leaves, and S is the number of samples.

• CoalescentRateLikelihood -- An experimental callable (not a Distribution) that provides a transposed version of CoalescentTimesWithRate, making rate_grid elements independent for compatibility with pyro.plate and poutine.markov. This is particularly useful for sequential Monte Carlo filtering.

Additionally, the helper function bio_phylo_to_times extracts coalescent summary statistics from a Bio.Phylo tree object.

Internally, the module uses a _Phylogeny named tuple to efficiently represent timing information (times, signs, lineages, binomial coefficients) and employs memoization via _weak_memoize to cache phylogeny computations.

Usage

Use CoalescentTimes when modeling phylogenetic trees with constant population size and you need both sampling and density evaluation. Use CoalescentTimesWithRate when you have a time-varying piecewise constant coalescent rate (e.g., from an epidemiological SIR model) and only need the likelihood. Use CoalescentRateLikelihood when you need to decompose the likelihood across time steps for use with pyro.plate or sequential inference methods like SMCFilter.

Code Reference

Source Location

• Repository: Pyro
• File: pyro/distributions/coalescent.py

Signature

class CoalescentTimes(TorchDistribution):
    arg_constraints = {"leaf_times": constraints.real, "rate": constraints.positive}
    def __init__(self, leaf_times, rate=1.0, *, validate_args=None)

class CoalescentTimesWithRate(TorchDistribution):
    arg_constraints = {"leaf_times": constraints.real, "rate_grid": constraints.positive}
    def __init__(self, leaf_times, rate_grid, *, validate_args=None)

class CoalescentRateLikelihood:
    def __init__(self, leaf_times, coal_times, duration, *, validate_args=None)
    def __call__(self, rate_grid, t=slice(None))

def bio_phylo_to_times(tree, *, get_time=None)

Import

from pyro.distributions import CoalescentTimes
from pyro.distributions import CoalescentTimesWithRate
from pyro.distributions.coalescent import CoalescentRateLikelihood, bio_phylo_to_times

I/O Contract

CoalescentTimes Inputs

Name	Type	Required	Description
leaf_times	torch.Tensor	Yes	Vector of times of sampling events (leaf nodes in the phylogeny). Can be arbitrary real numbers with arbitrary order and duplicates.
rate	torch.Tensor or float	No	Base coalescent rate (pairwise rate of coalescence) under a constant population size model. Defaults to 1.0.
validate_args	bool or None	No	Whether to enable argument validation.

CoalescentTimes Outputs

Name	Type	Description
sample	torch.Tensor	Sorted set of binary coalescent times with shape batch_shape + (leaf_times.size(-1) - 1,).
log_prob	torch.Tensor	Log probability density of the coalescent times, with shape batch_shape.

CoalescentTimesWithRate Inputs

Name	Type	Required	Description
leaf_times	torch.Tensor	Yes	Tensor of leaf sampling times.
rate_grid	torch.Tensor	Yes	Tensor of piecewise constant base coalescent rates. The rightmost dimension is time, representing intervals (-inf,1], [1,2], ..., [T-1,inf).

CoalescentTimesWithRate Outputs

Name	Type	Description
log_prob	torch.Tensor	leaf_times, rate_grid).

CoalescentRateLikelihood Inputs

Name	Type	Required	Description
leaf_times	torch.Tensor	Yes	Tensor of leaf sampling times.
coal_times	torch.Tensor	Yes	Sorted tensor of coalescent times with coal_times.size(-1) = leaf_times.size(-1) - 1.
duration	int	Yes	Size of the rate grid (must be >= 2).

CoalescentRateLikelihood Call Outputs

Name	Type	Description
likelihood	torch.Tensor	Log-likelihood contribution for the given rate_grid (or a single time slice).

Usage Examples

Basic CoalescentTimes

import torch
from pyro.distributions import CoalescentTimes

# Define leaf sampling times (e.g., 5 leaves)
leaf_times = torch.tensor([0.0, 0.1, 0.3, 0.5, 0.8])

# Create distribution with default rate=1.0
dist = CoalescentTimes(leaf_times)

# Sample coalescent times (4 events for 5 leaves)
coal_times = dist.sample()
print(coal_times.shape)  # torch.Size([4])

# Compute log probability
log_p = dist.log_prob(coal_times)

CoalescentTimesWithRate for Piecewise Constant Rate

import torch
from pyro.distributions import CoalescentTimesWithRate

leaf_times = torch.tensor([0.0, 1.0, 2.0, 3.0])
rate_grid = torch.tensor([1.0, 2.0, 1.5, 0.8, 1.2])

dist = CoalescentTimesWithRate(leaf_times, rate_grid)
# Compute log_prob given observed coalescent times
coal_times = torch.tensor([0.5, 1.5, 2.5])
log_p = dist.log_prob(coal_times)

CoalescentRateLikelihood with pyro.plate

import pyro
import torch
from pyro.distributions.coalescent import CoalescentRateLikelihood

leaf_times = torch.tensor([0.0, 1.0, 2.0, 3.0])
coal_times = torch.tensor([0.5, 1.5, 2.5])
rate_grid = torch.ones(5)

likelihood = CoalescentRateLikelihood(leaf_times, coal_times, len(rate_grid))

# Vectorized version using pyro.plate
with pyro.plate("time", len(rate_grid)):
    pyro.factor("coalescent", likelihood(rate_grid))

Extracting Times from Bio.Phylo Tree

from pyro.distributions.coalescent import bio_phylo_to_times

# Given a Bio.Phylo tree object
leaf_times, coal_times = bio_phylo_to_times(tree)

Related Pages

• Principle:Pyro_ppl_Pyro_Coalescent_Process