From 8face72241fca16414c4ac81fcf89a0d0017c6ad Mon Sep 17 00:00:00 2001
From: Savita Karthikeyan <savitakartik@gmail.com>
Date: Wed, 23 Aug 2023 09:11:52 +0100
Subject: [PATCH] first attempt at ancestor span heatmap

---
 model.py       | 32 ++++++++++++++++++++++++++++++++
 pages/nodes.py | 11 +++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/model.py b/model.py
index e53d019..a5d05af 100644
--- a/model.py
+++ b/model.py
@@ -479,3 +479,35 @@ 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, win_size_y=1):
+        edges_df = self.edges_df
+        num_x = int(np.ceil(edges_df.right.max() / win_size_x))
+        num_y = int(np.ceil(edges_df.child_time.max() / win_size_y))
+        anc_spans=[]
+        x_start = 0 # redo start to ts.start, set span to min and max of win/span
+        x=[]
+        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.append(x_end)
+                y.append(y_start)
+                anc_spans.append(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)].span.mean())
+        df = pd.DataFrame(
+            {
+                "Genomic position": x,
+                "Time": y,
+                "Mean ancestor span": anc_spans,
+            }
+        )
+        return df
\ No newline at end of file
diff --git a/pages/nodes.py b/pages/nodes.py
index 46020f0..83c4750 100644
--- a/pages/nodes.py
+++ b/pages/nodes.py
@@ -10,7 +10,7 @@
 from plot_helpers import filter_points
 from plot_helpers import hover_points
 from plot_helpers import make_hist_matplotlib
-
+import model
 
 def page(tsm):
     hv.extension("matplotlib")
@@ -60,4 +60,11 @@ 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=100_000, win_size_y=5_000)
+    #nc_span_data = tsm.calc_anc_spans(win_size_x=1, win_size_y=1)
+    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)