Dev

batch.py

@@ -533,7 +533,8 @@ def process_feed(input_args, input_cur, callback=None):
                 'max_ndiffs': max(ndiffs),
                 'mean_ndiffs': sum(ndiffs) / len(ndiffs),
                 'mutations': mutations[lineage],
-                'infections': infection_prediction[lineage],
+                'infections': infection_prediction[lineage]['infections'],
+                'hunepi': infection_prediction[lineage]['hunepi'] if 'hunepi' in infection_prediction[lineage] else None,
                 'raw_lineage': lname
             }
 

covizu/utils/batch_utils.py

@@ -277,8 +277,17 @@ def find_ne(tree, labels_filename):
         }
 
         #Run skyline estimation
-        alpha = betacoal.maxlik(tree)
-        skyline = (skyline.multi.phylo(tree, alpha$p1))
+        if (is.binary(tree)) {
+            alpha <- list(pi=1.999)
+        } else {
+            # alpha = betacoal.maxlik(tree)
+            Inter <- coalescent.intervals.multi(tree)
+            optimx::optimx(par = 1.5, fn = function(x) 
+              -skyline.multi.coalescentIntervals(Inter, x, epsilon=0)$logL, 
+              lower = 0.001, upper = 1.999, method = "L-BFGS-B")
+        }
+
+        skyline <- skyline.multi.phylo(tree, alpha$p1)
 
         #Output skyline estimation
         pop_sizes <- head(skyline$population.size, n = 5)
@@ -478,7 +487,7 @@ def make_beadplots(
                     beaddict['nodes'][variant].append(
                         [coldate, accn, location, label1])
 
-                inf_predict.update({lineage: 0})
+                inf_predict.update({lineage: {'infections': 0}})
             else:
                 # generate beadplot data
                 ctree = clustering.consensus(
@@ -542,7 +551,12 @@ def make_beadplots(
                 ''')
 
                 predicted_infections = list(robjects.r('predicted_infections'))[0]
-                inf_predict.update({lineage: predicted_infections})
+                inf_predict.update({lineage: {'infections': predicted_infections, 
+                                              'hunepi': 
+                                                {'h': summary_stats['shannons_diversity'], 
+                                                 'u': summary_stats['unsampled_lineage_count'], 
+                                                 'ne': summary_stats['Ne'], 
+                                                 'pi': summary_stats['pi']}}})
 
                 ctree = beadplot.annotate_tree(ctree, label_dict)
                 beaddict = beadplot.serialize_tree(ctree)

js/covizu.js

@@ -1019,7 +1019,7 @@ function export_csv() {
 
   // write lineage-level information to CSV file for download
   var csvFile = 'lineage,mean.diffs,clock.residual,num.cases,num.variants,min.coldate,max.coldate,' +
-      'mean.coldate,pred.cases';
+      'mean.coldate,pred.cases,H,U,Ne,Pi';
   var lineage_info = []
   for (tip of all_tips) {
     if (tip.isTip === undefined || tip.isTip)
@@ -1032,7 +1032,11 @@ function export_csv() {
         formatDate(tip.first_date),
         formatDate(tip.last_date),
         formatDate(tip.mcoldate),
-        Math.round(tip.infections)
+        Math.round(tip.infections),
+        tip.hunepi.h,
+        tip.hunepi.u,
+        tip.hunepi.ne,
+        tip.hunepi.pi
        ]);
   }
   csvFile = csvFile + "\n" + lineage_info.join("\n");

server/parseCluster.js

@@ -366,6 +366,7 @@ const map_tips = (cidx, labels, root, tips, tip_labels, cluster) => {
   tips[root_idx].nsamples = tip_stats.nsamples;
   tips[root_idx].mutations = tip_stats.mutations;
   tips[root_idx].infections = tip_stats.infections;
+  tips[root_idx].hunepi = tip_stats.hunepi;
 
   // calculate residual from mean differences and mean collection date - fixes #241
   let times = coldates.map(x => utcDate(x).getTime()),