@@ -595,43 +595,57 @@ def compute_ancestor_spans_heatmap_data(self, num_x_bins, num_y_bins):
595
595
"""
596
596
Calculates the average ancestor span in a genomic-time window
597
597
"""
598
- nodes_df = self .nodes_df [self .nodes_df .ancestors_span != - np .inf ]
599
- nodes_df = nodes_df .reset_index (drop = True )
600
- nodes_left = nodes_df .child_left
601
- nodes_right = nodes_df .child_right
602
- nodes_time = nodes_df .time
603
-
604
- x_bins = np .linspace (nodes_left .min (), nodes_right .max (), num_x_bins + 1 )
605
- y_bins = np .linspace (0 , nodes_time .max (), num_y_bins + 1 )
606
- heatmap_counts = np .zeros ((num_x_bins , num_y_bins ))
607
-
608
- x_starts = np .digitize (nodes_left , x_bins , right = True )
609
- x_ends = np .digitize (nodes_right , x_bins , right = True )
610
- y_starts = np .digitize (nodes_time , y_bins , right = True )
611
-
612
- for u in range (len (nodes_left )):
613
- x_start = max (0 , x_starts [u ] - 1 )
614
- x_end = max (0 , x_ends [u ] - 1 )
615
- y_bin = max (0 , y_starts [u ] - 1 )
616
- heatmap_counts [x_start : x_end + 1 , y_bin ] += 1
617
-
618
- x_coords = np .repeat (x_bins [:- 1 ], num_y_bins )
619
- y_coords = np .tile (y_bins [:- 1 ], num_x_bins )
620
- overlapping_node_count = heatmap_counts .flatten ()
621
- overlapping_node_count [overlapping_node_count == 0 ] = 1
622
- # FIXME - better way to avoid log 0 above?
623
- df = pd .DataFrame (
624
- {
625
- "position" : x_coords .flatten (),
626
- "time" : y_coords .flatten (),
627
- "overlapping_node_count_log10" : np .log10 (overlapping_node_count ),
628
- "overlapping_node_count" : overlapping_node_count ,
629
- }
630
- )
631
- return df .astype (
632
- {
633
- "position" : "int" ,
634
- "time" : "int" ,
635
- "overlapping_node_count" : "int" ,
636
- }
637
- )
598
+ if self .ts .time_units == tskit .TIME_UNITS_UNCALIBRATED :
599
+ logger .warning (
600
+ "Cannot compute ancestor spans for uncalibrated tree sequence"
601
+ )
602
+ return pd .DataFrame (
603
+ {
604
+ "position" : [],
605
+ "time" : [],
606
+ "overlapping_node_count_log10" : [],
607
+ "overlapping_node_count" : [],
608
+ }
609
+ )
610
+ else :
611
+ nodes_df = self .nodes_df [self .nodes_df .ancestors_span != - np .inf ]
612
+ nodes_df = nodes_df .reset_index (drop = True )
613
+ nodes_left = nodes_df .child_left
614
+ nodes_right = nodes_df .child_right
615
+ nodes_time = nodes_df .time
616
+
617
+ x_bins = np .linspace (nodes_left .min (), nodes_right .max (), num_x_bins + 1 )
618
+ y_bins = np .linspace (0 , nodes_time .max (), num_y_bins + 1 )
619
+ heatmap_counts = np .zeros ((num_x_bins , num_y_bins ))
620
+
621
+ x_starts = np .digitize (nodes_left , x_bins , right = True )
622
+ x_ends = np .digitize (nodes_right , x_bins , right = True )
623
+ y_starts = np .digitize (nodes_time , y_bins , right = True )
624
+
625
+ for u in range (len (nodes_left )):
626
+ x_start = max (0 , x_starts [u ] - 1 )
627
+ x_end = max (0 , x_ends [u ] - 1 )
628
+ y_bin = max (0 , y_starts [u ] - 1 )
629
+ heatmap_counts [x_start : x_end + 1 , y_bin ] += 1
630
+
631
+ x_coords = np .repeat (x_bins [:- 1 ], num_y_bins )
632
+ y_coords = np .tile (y_bins [:- 1 ], num_x_bins )
633
+ overlapping_node_count = heatmap_counts .flatten ()
634
+ overlapping_node_count [overlapping_node_count == 0 ] = 1
635
+ # FIXME - better way to avoid log 0 above?
636
+ df = pd .DataFrame (
637
+ {
638
+ "position" : x_coords .flatten (),
639
+ "time" : y_coords .flatten (),
640
+ "overlapping_node_count_log10" : np .log10 (overlapping_node_count ),
641
+ "overlapping_node_count" : overlapping_node_count ,
642
+ }
643
+ )
644
+ return df .astype (
645
+ {
646
+ "position" : "int" ,
647
+ "time" : "int" ,
648
+ "overlapping_node_count_log10" : "int" ,
649
+ "overlapping_node_count" : "int" ,
650
+ }
651
+ )
0 commit comments