Local Branching Index

src/actions/colors.js

@@ -30,7 +30,7 @@ export const changeColorBy = (providedColorBy = undefined) => { // eslint-disabl
     /* step 1: calculate the required colour scale */
     const version = controls.colorScale === undefined ? 1 : controls.colorScale.version + 1;
     // console.log("updateColorScale setting colorScale to ", version);
-    const colorScale = getColorScale(colorBy, tree, controls.geneLength, metadata.colorOptions, version);
+    const colorScale = getColorScale(colorBy, tree, controls.geneLength, metadata.colorOptions, version, controls.absoluteDateMaxNumeric);
     /*   */
     if (colorBy.slice(0, 3) === "gt-" && controls.geneLength) {
       colorScale.genotype = parseGenotype(colorBy, controls.geneLength);

src/components/tree/phyloTree.js

@@ -1,3 +1,4 @@
+/* eslint-disable */
 import _debounce from "lodash/debounce";
 import { event } from "d3-selection";
 import { min, max, sum } from "d3-array";

src/util/getColorScale.js

@@ -5,6 +5,7 @@ import { interpolateHcl } from "d3-interpolate";
 import { genericDomain, colors, genotypeColors, reallySmallNumber, reallyBigNumber } from "./globals";
 import { parseGenotype } from "./getGenotype";
 import { getAllValuesAndCountsOfTraitsFromTree } from "./treeTraversals";
+import { setLBI } from "./localBranchingIndex";
 
 /**
 * what values (for colorBy) are present in the tree and not in the color_map?
@@ -60,7 +61,7 @@ const genericScale = (cmin, cmax, vals = false) => {
 
 
 const minMaxAttributeScale = (nodes, attr, options) => {
-  if (options.vmin && options.vmax) {
+  if (Object.prototype.hasOwnProperty.call(options, "vmin") && Object.prototype.hasOwnProperty.call(options, "vmax")) {
     return genericScale(options.vmin, options.vmax);
   }
   const vals = nodes.map((n) => n.attr[attr])
@@ -112,9 +113,11 @@ const discreteAttributeScale = (nodes, attr) => {
     .range(colorList);
 };
 
-const getColorScale = (colorBy, tree, geneLength, colorOptions, version) => {
+const getColorScale = (colorBy, tree, geneLength, colorOptions, version, absoluteDateMaxNumeric) => {
   let colorScale;
   let continuous = false;
+  let error = false;
+
   if (!tree.nodes) {
     // make a dummy color scale before the tree is in place
     continuous = true;
@@ -136,6 +139,16 @@ const getColorScale = (colorBy, tree, geneLength, colorOptions, version) => {
         colorScale = scaleOrdinal().domain(domain).range(genotypeColors);
       }
     }
+  } else if (colorBy === "lbi") {
+    try {
+      setLBI(tree.nodes, absoluteDateMaxNumeric, colorOptions.lbi.tau, colorOptions.lbi.timeWindow);
+      // colorScale = minMaxAttributeScale(tree.nodes, "lbi", colorOptions.lbi); /* colour ramp over all values */
+      colorScale = minMaxAttributeScale(undefined, undefined, {vmin: 0, vmax: 0.7}); /* ramp over [0, 0.7] like nextflu */
+      continuous = true;
+    } catch (e) {
+      console.error("Setting LBI failed.", e);
+      error = true;
+    }
   } else if (colorOptions && colorOptions[colorBy]) {
     if (colorOptions[colorBy].color_map) {
       // console.log("Sweet - we've got a color_map for ", colorBy)
@@ -166,11 +179,15 @@ const getColorScale = (colorBy, tree, geneLength, colorOptions, version) => {
       colorScale = minMaxAttributeScale(tree.nodes, colorBy, colorOptions[colorBy]);
     }
   } else {
-    // This shouldn't ever happen!
-    // console.log("no colorOptions for ", colorBy, " returning minMaxAttributeScale")
+    error = true;
+  }
+
+  if (error) {
+    console.error("no colorOptions for ", colorBy, " returning minMaxAttributeScale");
     continuous = true;
     colorScale = minMaxAttributeScale(tree.nodes, colorBy, colorOptions[colorBy]);
   }
+
   return {
     scale: colorScale,
     continuous: continuous,

src/util/localBranchingIndex.js

@@ -0,0 +1,99 @@
+
+/* sets each node in the tree to alive=true if it has at least one descendent with current=true */
+const setNodeAlive = (node, cutoff) => {
+  if (node.children) {
+    let aliveChildren = false;
+    for (let i = 0, c = node.children.length; i < c; i++) {
+      setNodeAlive(node.children[i], cutoff);
+      aliveChildren = aliveChildren || node.children[i].alive;
+    }
+    node.alive = aliveChildren;
+  } else {
+    node.alive = node.attr.num_date > cutoff;
+  }
+};
+
+/* for each node, calculate the exponentially attenuated tree length below the node
+the polarizer is send "up", i.e. to parents */
+function calcUpPolarizers(node, LBItau) {
+  node.up_polarizer = 0;
+  if (node.children) {
+    for (let i = 0; i < node.children.length; i++) {
+      calcUpPolarizers(node.children[i], LBItau);
+      node.up_polarizer += node.children[i].up_polarizer;
+    }
+  }
+  const bl = node.clock_length / LBItau;
+  node.up_polarizer *= Math.exp(-bl);
+  if (node.alive) { // only alive branches contribute anything
+    node.up_polarizer += LBItau * (1 - Math.exp(-bl));
+  }
+}
+
+/* for each node, calculate the exponentially attenuated tree length above the node,
+that is "outside" the clade defined by this node. this down polarizer is send to children */
+function calcDownPolarizers(node, LBItau) {
+  if (node.children) {
+    for (let i1 = 0; i1 < node.children.length; i1++) {
+      node.children[i1].down_polarizer = node.down_polarizer;
+      for (let i2 = 0; i2 < node.children.length; i2++) {
+        if (i1 !== i2) {
+          node.children[i1].down_polarizer += node.children[i2].up_polarizer;
+        }
+      }
+      // account for the attenuation over the branch_length
+      const bl = node.children[i1].clock_length / LBItau;
+      node.children[i1].down_polarizer *= Math.exp(-bl);
+      if (node.children[i1].alive) { // the branch contributes only when the node is alive
+        node.children[i1].down_polarizer += LBItau * (1 - Math.exp(-bl));
+      }
+      calcDownPolarizers(node.children[i1], LBItau);
+    }
+  }
+}
+
+function calcPolarizers(node, LBItau) {
+  calcUpPolarizers(node, LBItau);
+  node.down_polarizer = 0; // set the down polarizer of the root to 0
+  calcDownPolarizers(node, LBItau);
+}
+
+/* calculate the LBI for all nodes downstream of node
+allnodes is provided for easy normalization at the end
+Side effects: adds the following to nodes:
+  n.alive
+  n.up_polarizer
+  n.down_polarizer
+  n.clock_length
+  n.attr.lbi
+Clealy n.attr.lbi is useful, but can we avoid storing those other values?
+*/
+export const setLBI = (nodes, maxDateInTree, LBItau, LBItimeWindow) => {
+  // console.time('LBI');
+  const LBIcutoff = maxDateInTree - LBItimeWindow;
+  nodes.forEach((d) => {
+    if (d.children) {
+      for (let i = 0; i < d.children.length; i++) {
+        d.children[i].clock_length = d.children[i].tvalue - d.tvalue;
+      }
+    }
+  });
+  nodes[0].clock_length = 0;
+  setNodeAlive(nodes[0], LBIcutoff);
+  calcPolarizers(nodes[0], LBItau);
+  let maxLBI = 0;
+  nodes.forEach((d) => {
+    d.attr.lbi = d.down_polarizer;
+    if (d.children) {
+      for (let i = 0; i < d.children.length; i++) {
+        d.attr.lbi += d.children[i].up_polarizer;
+      }
+    }
+    if (d.attr.lbi > maxLBI) {
+      maxLBI = d.attr.lbi;
+    }
+  });
+  // normalize the LBI to range [0,1]
+  nodes.forEach((d) => {d.attr.lbi /= maxLBI;});
+  // console.timeEnd('LBI');
+};