Time-Stratified SFS

demestats.tsfs extends the expected SFS calculation to a fixed time grid. Instead of returning one expected count for each derived-allele configuration, it partitions that expected count by mutation age.

Given a grid

0 = t[0] < t[1] < ... < t[M] = inf,

the entry tsfs[i, ...] is the expected number of mutations that:

  • arose during the time interval [t[i], t[i + 1]), and

  • subtend the corresponding lineage configuration on the remaining axes.

Summing over the first axis recovers the ordinary expected SFS.

Basic usage

import jax.numpy as jnp
from demestats.sfs import ExpectedSFS
from demestats.tsfs import ExpectedTSFS

time_grid = [0.0, 50.0, 200.0, jnp.inf]

etsfs = ExpectedTSFS(
    demo,
    num_samples={"A": 4, "B": 4},
    time_grid=time_grid,
)

tsfs = etsfs()
print(tsfs.shape)
# (3, 5, 5)

esfs = ExpectedSFS(demo, num_samples={"A": 4, "B": 4})()
print(jnp.allclose(tsfs.sum(axis=0), esfs))
# True

The first axis indexes time bins. The remaining axes use the same convention as ExpectedSFS: for a deme with n haploid samples, the corresponding axis has length n + 1.

Time grid requirements

time_grid is fixed when the ExpectedTSFS object is created. It must:

  • be one-dimensional,

  • be strictly increasing,

  • start at 0, and

  • end at inf.

The grid is interpreted in the same time units as the input demes.Graph. Internal rescaling is handled automatically, so you should always specify the grid in the original demographic-model units.

Output interpretation

If num_samples={"A": nA, "B": nB}, then:

  • tsfs.shape == (M, nA + 1, nB + 1),

  • tsfs[i, j, k] is the expected number of mutations from time bin i with j derived alleles in A and k derived alleles in B,

  • tsfs.sum(axis=0) is the ordinary expected SFS, and

  • tsfs[:, 0, 0] and tsfs[:, -1, -1] are zero, just as in the ordinary SFS.

The total expected number of mutations in a time interval is tsfs[i].sum().

Parameter overrides

ExpectedTSFS accepts the same parameter override dictionary as ExpectedSFS. This makes it easy to evaluate the TSFS across different demographic parameter settings without rebuilding the object:

etsfs = ExpectedTSFS(demo, num_samples={"A": 4}, time_grid=[0.0, 100.0, jnp.inf])

params = {
    etsfs.variable_for(("demes", 0, "epochs", 0, "start_size")): 2_000.0,
}

tsfs = etsfs(params)

Pruning

ExpectedTSFS supports the same prune= option as ExpectedSFS. Pruning only changes the internal dynamic-programming state. It does not change the output shape.

from demestats.tsfs import ExpectedTSFS

etsfs = ExpectedTSFS(
    demo,
    num_samples={"A": 12, "B": 8},
    time_grid=[0.0, 100.0, 500.0, jnp.inf],
    prune={"A": 4},
)

For the pruning syntax and caveats, see Pruning.

Current scope

The first version of demestats.tsfs focuses on direct TSFS evaluation through ExpectedTSFS(...)(params). The ExpectedSFS.tensor_prod helper does not yet have a TSFS counterpart.

For API details, see the generated reference under API.