Ancestor span heatmap

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

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)