diff --git a/src/actions/colors.js b/src/actions/colors.js index 3a17144c2..5aae1bec4 100644 --- a/src/actions/colors.js +++ b/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); diff --git a/src/components/tree/phyloTree.js b/src/components/tree/phyloTree.js index 1807fc67e..88c05124e 100644 --- a/src/components/tree/phyloTree.js +++ b/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"; diff --git a/src/util/getColorScale.js b/src/util/getColorScale.js index acb5804c9..efe7385e7 100644 --- a/src/util/getColorScale.js +++ b/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, diff --git a/src/util/localBranchingIndex.js b/src/util/localBranchingIndex.js new file mode 100644 index 000000000..86b27e598 --- /dev/null +++ b/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'); +};