adding sample node-related information as columns to the mutations dataframe

[savitakartik]

It would be useful to bring in external information about the sample nodes to tsqc. For example, we could show a bar chart of sample node frequencies for each group of self-reported ethnicities under a mutation. For this, we should have a numba function which takes in the sample information (ethnicities in this case) and return a dataframe with the shape, number of mutations x number of information groups. Then, we could add these columns to the mutations_df when we create it. This is ideal as Holoviews would then pass the entire dataframe to the browser, making it quicker to interact with individual mutation points, rather than calculating node frequencies on the fly.

[jeromekelleher]

The simplest thing we can do here is to compute frequencies as defined by population. Assuming all sample nodes are assigned to a population (in the population table) then we can compute the frequency of every node in the tree sequence like:

    parent = np.zeros(num_nodes, dtype=np.int32) - 1
    num_pop_samples = np.zeros((num_nodes, num_populations), dtype=np.int32)
    for pop in range(num_populations):
           samples =  np.logical_and(nodes_population == pop, nodes_flags == 1) # Not quite right!
           num_pop_samples[samples, pop] = 1

    while tree_pos.next():
        for j in range(tree_pos.out_range[0], tree_pos.out_range[1]):
            e = tree_pos.edge_removal_order[j]
            c = edges_child[e]
            p = edges_parent[e]
            parent[c] = -1
            u = p
            while u != -1:
                num_pop_samples[u] -= num_pop_samples[c]
                u = parent[u]

        for j in range(tree_pos.in_range[0], tree_pos.in_range[1]):
            e = tree_pos.edge_insertion_order[j]
            p = edges_parent[e]
            c = edges_child[e]
            parent[c] = p
            u = p
            while u != -1:
                num_pop_samples[u] += num_pop_samples[c]
                u = parent[u]

Then, we can use num_pop_samples to compute the frequency of every mutation within each population in this function: https://github.com/tskit-dev/tsqc/blob/26d3ee9a6a55948a1474e5c634862f9d0e3d6c33/tsqc/model.py#L179

[savitakartik]

Thanks, @jeromekelleher .

For the UI part, here's what I have so far. I'm starting with a muts_df with num_pop additional columns containing population frequencies for each mutation.

import pandas as pd
import holoviews as hv
from holoviews.streams import Selection1D
from holoviews import opts

hv.extension('bokeh')

# Retrieve frequency data for a given mutation id in the format required for hv.Bars
def gen_freqs(mut_id):
    return pd.DataFrame({
        'pop': ['pop_A', 'pop_B', 'pop_C'],
        'freq': muts_df.loc[mut_id, ['pop_A', 'pop_B', 'pop_C']].values
    })

# Create an NdOverlay of Bars for all the mutations
pop_data = hv.NdOverlay({i: hv.Bars(gen_freqs(i), 'pop', 'freq')
                         for i in range(100)})

# Create a scatter plot
points = hv.Scatter(muts_df, kdims=['position'], vdims=[y_dim, 'id', 'inheritors', 'descendants'])

# Link a Selection1D stream to the scatter plot
stream = Selection1D(source=points)

# Update the bar plot based on the selected point
def plot_pop_freqs(index):
    if not index:
        # Return an empty Bars and an empty Table when nothing is selected
        empty_bars = hv.Bars([], kdims='pop', vdims='freq').opts(width=200, height=200, title='Tap on a mutation')
        return empty_bars
    selected_index = index[0]
    bars = pop_data[selected_index].opts(width=200, height=200, framewise=True, title=f'Mutation {selected_index}', ylim=(0, 20))
    return bars

#Update the data table based on the selected point
def make_mutation_table(index):
    if not index:
        return hv.Table([], kdims='Details', vdims='Value').opts(width=200, height=200)
    selected_index = index[0]
    num_inheritors = muts_df.loc[selected_index, 'inheritors']
    num_descendants = muts_df.loc[selected_index, 'descendants']
    table_data = [('Num Inheritors', num_inheritors), ('Num Descendants', num_descendants)]
    table = hv.Table(table_data, kdims='Details', vdims='Value').opts(width=200, height=200)
    return table

pop_freq_plot = hv.DynamicMap(plot_pop_freqs, streams=[stream])
mutation_table = hv.DynamicMap(make_mutation_table, streams=[stream])

layout = (points.opts(size=10, tools=['tap', 'hover']) + pop_freq_plot
            + mutation_table)

layout.opts(opts.Scatter(width=600, height=400, size=10, framewise=True))

[savitakartik]

population_frequencies_v0.mp4

[savitakartik]

This was just the first attempt that worked. There are issues, for example, I don't think this will work for large numbers of mutations:

pop_data = hv.NdOverlay({i: hv.Bars(gen_freqs(i), 'pop', 'freq')
                         for i in range(100)})

If you think this is useful, I'll try and add this part to tsqc. Will likely need help with creating the dataframe though - would be great to discuss next week if you are around.

[jeromekelleher]

This looks great! Let's get some version of it in (fake data would be fine initially)

I'll help with computing the dataframe next week, no problem

[savitakartik]

fixed with PRs #141 #145