diff --git a/model.py b/model.py index e53d019..a72380f 100644 --- a/model.py +++ b/model.py @@ -479,3 +479,43 @@ def calc_mutations_per_tree(self): mutations_per_tree = np.zeros(self.ts.num_trees, dtype=np.int64) mutations_per_tree[unique_values] = counts return mutations_per_tree + + def calc_anc_spans(self, win_size_x=1_000_000, win_size_y=5_000): + edges_df = self.edges_df + num_x = int(np.ceil(edges_df.right.max() - edges_df.right.min()) / win_size_x) + num_y = int(np.ceil(edges_df.child_time.max() / win_size_y)) + anc_spans = np.zeros((num_x, num_y)) + x_start = edges_df.left.min() + x = np.zeros((num_x, num_y)) + y = np.zeros((num_x, num_y)) + for i in range(num_x): + x_start = i * win_size_x + x_end = x_start + win_size_x + + for j in range(num_y): + y_start = j * win_size_y + y_end = y_start + win_size_y + x[i][j] = x_end + y[i][j] = y_start + tmp_df = edges_df[ + ( + ((x_start >= edges_df.left) & (x_start < edges_df.right)) + | ((x_end > edges_df.left) & (x_end <= edges_df.right)) + ) + & (edges_df.child_time >= y_start) + & (edges_df.child_time < y_end) + ] + n = np.unique(tmp_df.child).shape[0] + anc_spans[i][j] = np.sum(np.clip(tmp_df.span, None, win_size_x)) / n + anc_spans = anc_spans.flatten() + x = x.flatten() + y = y.flatten() + + df = pd.DataFrame( + { + "Genomic_position": x, + "Time": y, + "Mean_ancestor_span": anc_spans, + } + ) + return df diff --git a/pages/nodes.py b/pages/nodes.py index 46020f0..3c4de4a 100644 --- a/pages/nodes.py +++ b/pages/nodes.py @@ -60,4 +60,13 @@ def page(tsm): pn.pane.Markdown("# Plot Options"), log_y_checkbox, ) - return pn.Column(main, hist_panel, plot_options) + + anc_span_data = tsm.calc_anc_spans(win_size_x=1_000_000, win_size_y=5_000) + heatmap = hv.HeatMap(anc_span_data).opts( + width=config.PLOT_WIDTH, + height=config.PLOT_HEIGHT, + tools=["hover"], + colorbar=True, + ) + + return pn.Column(main, hist_panel, heatmap, plot_options)