set first and last dates for unsampled, #58

covizu/beadplot.py

@@ -1,7 +1,6 @@
 import sys
 import json
 import argparse
-from typing import OrderedDict
 from covizu.clustering import consensus
 from Bio import Phylo
 from io import StringIO
@@ -146,25 +145,39 @@ def serialize_tree(tree):
             intermed.sort()  # ISO dates sort in increasing order
             variant = intermed[0][1]  # use accession of earliest sample to ID variant
 
+            # Map the nodes to length and ascending dates
+            first_date = intermed[0][0]
+            last_date = intermed[len(intermed) - 1][0]
+
             # populate list with samples
             obj['nodes'].update({variant: intermed})
         else:
             variant = 'unsampled'+str(us_count)
             obj['nodes'].update({variant: []})
             us_count += 1
 
+            # Set first date to first date of parent; last date to the earliest first date of a descendant
+            id = node
+            while(not parents[id].labels):
+                id = parents[id]
+
+            first_date = obj['nodes'][variant_d[parents[id]]][0][0]
+            last_date = datetime.today().strftime('%Y-%m-%d')
+            for child in node:
+                if(child.labels):
+                    intermed = [label.split('|')[::-1] for label in child.labels]
+                    intermed.sort()
+                    if (intermed[0][0] < last_date):
+                        last_date = intermed[0][0]
+
         variant_d.update({node: variant})
 
         if node is tree.root:
             continue  # no edge
         parent = parents[node]
-
-        # Map the nodes to length and ascending dates
-        first_date = intermed[0][0]
-        last_date = intermed[len(intermed) - 1][0]
         date_diff = datetime.strptime(last_date, "%Y-%m-%d") - datetime.strptime(first_date, "%Y-%m-%d")
 
-        # prent ID, child ID, branch length, node support
+        # date width, first date, parent ID, child ID, branch length, node support
         if(variant_d[parent] in map_nodes):
             map_nodes[variant_d[parent]].append([date_diff, first_date, variant_d[parent], variant, round(node.branch_length, 2),
                              node.confidence])