Cumulative edge area (i.e. mutational area) over time

[hyanwong]

I wanted to calculate the mutational area over time of a large tree sequence, and compare that with the actual observed accumulation of mutations. It wasn't obvious the me how to get the cumulative area of the edges as a function of time, but a bit of AI querying got me the following fast implementation, so I'm sticking it here for posterity:

def cumulative_area(ts):
    """
    Calculate the cumulative area of edges, going from oldest to youngest
    (unique) node times in the tree sequence. This is equivalent to the
    area available for mutations, so is sometimes known as "mutational area"
    of the tree sequence. The last value in the returned `cumulative_times`
    vector is the total edge area in the tree sequence.

    Note that if there are any mutations above local roots, these do
    not occur on edges, and hence are not relevant to this function.
    
    Parameters:
    ts: tree_sequence object
    
    Returns:
    tuple: (times, cumulative_areas)
        times: array of unique node times where edges start or end
        cumulative_areas: cumulative edge area from an indefinitely long time ago to
            each timepoint in the `times` array
    """
    spans = ts.edges_right-ts.edges_left
    # We want to accumulate area from oldest parent to youngest child. To do this
    # while allowing np.unique to sort the times from old->young, we negate the times
    starts = -ts.nodes_time[ts.edges_parent]
    ends = -ts.nodes_time[ts.edges_child]
    
    # Create event arrays: each event has a time and a span change
    # Start events add span, end events subtract spans
    event_times = np.concatenate([starts, ends])
    event_spans = np.concatenate([spans, -spans])
    
    # Sort events by time
    sort_idx = np.argsort(event_times)
    event_times = event_times[sort_idx]
    event_spans = event_spans[sort_idx]

    # Combine duplicate timepoints by summing their span changes
    unique_times, indices, counts = np.unique(event_times, return_index=True, return_counts=True)
    span_changes = np.zeros_like(unique_times, dtype=float)


    # Accumulate span changes at identical timepoints. Spans have been sorted already
    for i, (idx, count) in enumerate(zip(indices, counts)):
        span_changes[i] = np.sum(event_spans[idx:idx+count])
        
    # Calculate running total of active spans
    active_spans = np.cumsum(span_changes)
    assert active_spans[-1] == 0  # at the last timepoint, we should have removed all edges
    active_spans = np.concatenate(([0.0], active_spans[:-1]))
                                  
    # Calculate time differences between consecutive positions
    # First element should be 0 since no area has been accumulated yet
    time_diffs = np.concatenate(([0.0], np.diff(unique_times)))
    
    # Current span at each position multiplied by time gives area increment
    area_increments = active_spans * time_diffs
    
    # Calculate cumulative area
    cumulative_areas = np.cumsum(area_increments)
    
    return -unique_times, cumulative_areas

Test to check that the eCDF of simulated mutation times agrees with the cumulative area

from matplotlib import pyplot as plt
import msprime
sample_sizes = [100, 1000, 10000, 100000]
fig, axes = plt.subplots(1, len(sample_sizes), figsize=(14, 4))
for ax, sample_size in zip(axes, sample_sizes):
    ts = msprime.sim_ancestry(sample_size, population_size=1000, sequence_length=1e5, recombination_rate=1e-7)
    ts = msprime.sim_mutations(ts, rate=1e-7)
    times, areas = cumulative_area(ts)
    print(
        "Total area from function:",
        areas[-1],
        "== total edge area directly calculated:",
        np.sum((ts.edges_right - ts.edges_left)*(ts.nodes_time[ts.edges_parent]-ts.nodes_time[ts.edges_child]))
    )
    ax.plot(times+1, areas/np.sum((ts.edges_right - ts.edges_left)*(ts.nodes_time[ts.edges_parent]-ts.nodes_time[ts.edges_child])))
    ax.set_xlim(ts.max_time, 1)
    ax.set_xscale("log")
    ax.set_title(f"Sample size={sample_size}")
    ax.ecdf(ts.mutations_time+1, compress=True, complementary=True, label="Mutations")

[petrelharp]

Nice; this seems like a good diagnostic plot. Just for fun, here's a version that avoids the for loop using np.bincount:

def cumulative_area2(ts):
    spans = ts.edges_right - ts.edges_left
    starts = -ts.nodes_time[ts.edges_parent]
    ends = -ts.nodes_time[ts.edges_child]
    event_times = np.concatenate([starts, ends])
    event_spans = np.concatenate([spans, -spans])
    times = np.unique(event_times)
    dt = np.diff(times)
    assert np.all(dt) > 0
    event_ind = np.searchsorted(times, event_times)
    cumulative_areas = np.cumsum(dt * np.cumsum(np.bincount(event_ind, weights=event_spans))[:-1])
    return -times, np.concatenate([[0.0], cumulative_areas])

[hyanwong]

Neat, thanks! much nicer.

[hyanwong]

Also seems that something like it might be a reasonable built-in tree sequence function

[petrelharp]

Once we add time_windows then we'll be able to get this from ts.segregating_sites(mode='branch', time_windows=...).

[hyanwong]

Huh, true! Good point.

[hyanwong]

A tree-based version can be obtained using the tree.edge_array indexes (although this is not incremental):

def cumulative_branch_length(tree):
    """
    Calculate the cumulative branch length in the tree going from oldest to youngest
    (unique) node times in the tree. This is equivalent to the
    area available for mutations. The last value in the returned `cumulative_lengths`
    vector is the total branch length of this tree.

    Note that if there are any mutations above local roots, these do
    not occur on edges, and hence are not relevant to this function.
    
    Parameters:
    tree: tskit.Tree object
    
    Returns:
    tuple: (times, cumulative_lengths)
        times: array of unique node times where branches start or end, sorted descending
        cumulative_lengths: cumulative branch lengths from an indefinitely long time ago to
            each timepoint in the `times` array
    """
    used_ids = tree.edge_array[tree.edge_array != tskit.NULL]
    starts = -ts.nodes_time[ts.edges_parent[used_ids]]
    ends = -ts.nodes_time[ts.edges_child[used_ids]]
    
    # Create event arrays: each event has a position and a num lineage change
    # Start events add lineages, end events subtract lineages
    event_times = np.concatenate([starts, ends])
    event_lineage_count = np.concatenate([np.ones(len(used_ids)), -np.ones(len(used_ids))])
    times = np.unique(event_times)
    dt = np.diff(times)
    assert np.all(dt) > 0
    event_ind = np.searchsorted(times, event_times)
    cumulative_areas = np.cumsum(dt * np.cumsum(np.bincount(event_ind, weights=event_lineage_count))[:-1])
    return -times, np.concatenate([[0.0], cumulative_areas])

[hyanwong]

With ts.segregating_sites(mode='branch', time_windows=..., windows="trees") this should also be possible incrementally (although it might blow up with large numbers of trees).