#1008 split tree and group into subtrees by state

src/actions/types.js

@@ -51,3 +51,4 @@ export const CHANGE_ZOOM = "CHANGE_ZOOM";
 export const SET_AVAILABLE = "SET_AVAILABLE";
 export const TOGGLE_SIDEBAR = "TOGGLE_SIDEBAR";
 export const TOGGLE_LEGEND = "TOGGLE_LEGEND";
+export const TOGGLE_SPLIT_TREE = "TOGGLE_SPLIT_TREE";

src/components/controls/choose-tree-split.js

@@ -0,0 +1,36 @@
+import React from "react";
+import Toggle from "./toggle";
+import { connect } from "react-redux";
+import { withTranslation } from "react-i18next";
+import { TOGGLE_SPLIT_TREE } from "../../actions/types";
+// todo: connect if needed
+
+@connect((state) => {
+    return {
+        colorBy: state.controls.colorBy,
+        splitTreeByTrait: state.tree.splitTreeByTrait
+    };
+})  
+/* Implements a button which splits the tree into visual trees per strain. */
+class ChooseTreeSplit extends React.Component {
+    render() {
+        const { t } = this.props;
+        return (
+            <div style={{margin: 5}}>
+                <Toggle
+                  display={true}
+                  on={this.props.splitTreeByTrait !== null}
+                  callback={() => 
+                        this.props.dispatch({type: TOGGLE_SPLIT_TREE, 
+                            // the presence of splitByTrait means it should be toggled off
+                            splitTreeByTrait: this.props.splitTreeByTrait !== null ? null : this.props.colorBy
+                        })}
+                  label={t("sidebar:Split tree by colored-by trait")}
+                />
+            </div>
+        );
+    }
+}
+
+const WithTranslation = withTranslation()(ChooseTreeSplit);
+export default WithTranslation;

src/components/controls/controls.js

@@ -8,6 +8,7 @@ import ChooseLayout from "./choose-layout";
 import ChooseDataset from "./choose-dataset";
 import ChooseSecondTree from "./choose-second-tree";
 import ChooseMetric from "./choose-metric";
+import ChooseTreeSplit from "./choose-tree-split";
 import PanelLayout from "./panel-layout";
 import GeoResolution from "./geo-resolution";
 import MapAnimationControls from "./map-animation";
@@ -36,6 +37,7 @@ function Controls({mapOn}) {
       <SidebarHeader>{t("sidebar:Tree Options")}</SidebarHeader>
       <ChooseLayout/>
       <ChooseMetric/>
+      <ChooseTreeSplit/>
       <ChooseBranchLabelling/>
       <SearchStrains/>
       <ChooseSecondTree/>

src/components/tree/phyloTree/change.js

@@ -1,6 +1,6 @@
 import { timerFlush } from "d3-timer";
 import { calcConfidenceWidth } from "./confidence";
-import { applyToChildren } from "./helpers";
+import { applyToChildren, setSplitTreeYValues, setYValues } from "./helpers";
 import { timerStart, timerEnd } from "../../../util/perf";
 import { NODE_VISIBLE } from "../../../util/globals";
 import { getBranchVisibility, strokeForBranch } from "./renderers";
@@ -275,14 +275,18 @@ export const change = function change({
   fill = undefined,
   visibility = undefined,
   tipRadii = undefined,
-  branchThickness = undefined
+  branchThickness = undefined,
+  resetTreeYValues = false,
+  splitTreeByTrait = null
 }) {
   // console.log("\n** phylotree.change() (time since last run:", Date.now() - this.timeLastRenderRequested, "ms) **\n\n");
   timerStart("phylotree.change()");
   const elemsToUpdate = new Set(); /* what needs updating? E.g. ".branch", ".tip" etc */
   const nodePropsToModify = {}; /* which properties (keys) on the nodes should be updated (before the SVG) */
   const svgPropsToUpdate = new Set(); /* which SVG properties shall be changed. E.g. "fill", "stroke" */
-  const useModifySVGInStages = newLayout; /* use modifySVGInStages rather than modifySVG. Not used often. */
+
+  /* use modifySVGInStages rather than modifySVG. Not used often. */
+  const useModifySVGInStages = newLayout || splitTreeByTrait !== null || resetTreeYValues; 
 
   /* calculate dt */
   const idealTransitionTime = 500;
@@ -318,7 +322,7 @@ export const change = function change({
     svgPropsToUpdate.add("stroke-width");
     nodePropsToModify["stroke-width"] = branchThickness;
   }
-  if (newDistance || newLayout || updateLayout || zoomIntoClade || svgHasChangedDimensions) {
+  if (newDistance || newLayout || updateLayout || zoomIntoClade || svgHasChangedDimensions || splitTreeByTrait) {
     elemsToUpdate.add(".tip").add(".branch.S").add(".branch.T").add(".branch");
     elemsToUpdate.add(".vaccineCross").add(".vaccineDottedLine").add(".conf");
     elemsToUpdate.add('.branchLabel').add('.tipLabel');
@@ -349,14 +353,31 @@ export const change = function change({
     this.nodes.forEach((d) => {d.update = true;});
   }
 
+  /* if the tree is being split by colored-by trait, update y values */
+  // todo: this is fairly slow: show something to show recalculating? or, get the d3 transition working?
+  if (splitTreeByTrait) {
+    setSplitTreeYValues(this.nodes, splitTreeByTrait);
+    this.params.isTreeSplitByTrait = true;
+  }
+  else if (resetTreeYValues)
+  {
+    setYValues(this.nodes);
+    this.params.isTreeSplitByTrait = false;
+  }
+
   /* run calculations as needed - these update properties on the phylotreeNodes (similar to updateNodesWithNewData) */
   /* distance */
   if (newDistance) this.setDistance(newDistance);
+
   /* layout (must run after distance) */
-  if (newDistance || newLayout || updateLayout) this.setLayout(newLayout || this.layout);
+  /* also used to split by traits */
+  if (newDistance || newLayout || updateLayout || splitTreeByTrait || resetTreeYValues)
+    this.setLayout(newLayout || this.layout);
+
   /* show confidences - set this param which actually adds the svg paths for
-     confidence intervals when mapToScreen() gets called below */
+    confidence intervals when mapToScreen() gets called below */
   if (showConfidences) this.params.confidence = true;
+
   /* mapToScreen */
   if (
     svgPropsToUpdate.has(["stroke-width"]) ||
@@ -365,7 +386,9 @@ export const change = function change({
     updateLayout ||
     zoomIntoClade ||
     svgHasChangedDimensions ||
-    showConfidences
+    showConfidences ||
+    splitTreeByTrait ||
+    resetTreeYValues
   ) {
     this.mapToScreen();
   }

src/components/tree/phyloTree/helpers.js

@@ -101,6 +101,73 @@ export const setYValuesRecursively = (node, yCounter) => {
  */
 export const setYValues = (nodes) => setYValuesRecursively(nodes[0], 0);
 
+// sorts all child nodes from a subtree into their subtrees
+const collectSubtrees = (startNode, subtree, trait, currentTraitValue, subtreeStack) => {
+  if (!startNode.children) return;
+
+  for (let i = startNode.children.length-1; i >= 0; i--) {
+    let currentNode = startNode.children[i];
+    let thisNodeTraitValue = getTraitFromNode(currentNode.n, trait);
+    // todo: nodes with no trait value? should they be grouped together?
+    if (thisNodeTraitValue !== currentTraitValue) { // doesn't belong on this subtree
+      let matchingSubtree = subtreeStack.find(s => s.traitValue === thisNodeTraitValue);
+      if (!matchingSubtree) {
+        subtreeStack.push({ subtreeNodes: [currentNode], traitValue: thisNodeTraitValue });
+        matchingSubtree = subtreeStack[subtreeStack.length-1];
+      }
+      else {
+        matchingSubtree.subtreeNodes.push(currentNode);
+      }
+      collectSubtrees(currentNode, matchingSubtree, trait, thisNodeTraitValue, subtreeStack);
+    }
+    else
+    {
+      // add this node to the subtree list, and continue
+      subtree.subtreeNodes.push(currentNode);
+      collectSubtrees(currentNode, subtree, trait, currentTraitValue, subtreeStack);
+    }    
+  }
+}
+
+/**
+ * setSplitTreeYValues - works similarly to setYValues above,
+ * but splits the tree by the given trait, grouping nodes with the
+ * same trait value together
+ * 
+ * todo: is 0 guaranteed to be root node?
+ */
+export const setSplitTreeYValues = (nodes, trait) => {
+  const subtreeStack = [{subtreeNodes: [nodes[0]], traitValue: getTraitFromNode(nodes[0].n, trait)}];  
+
+  // collect all the subtrees for a given trait, and group them together
+  collectSubtrees(nodes[0], subtreeStack[0], trait, subtreeStack[0].traitValue, subtreeStack);  
+  subtreeStack.sort((a, b) => {
+    if (a.traitValue < b.traitValue) return -1;
+    if (a.traitValue > b.traitValue) return 1;
+    return 0;
+  });
+
+  // todo: sort the subtrees by num_date
+  subtreeStack.sort((a, b) => {
+    if (a.traitValue == b.traitValue)
+      return 0;
+
+    let aDate = new Date(getTraitFromNode(a.subtreeNodes[0].n, "num_date"));
+    let bDate = new Date(getTraitFromNode(b.subtreeNodes[0].n, "num_date"));
+    if (aDate < bDate) return -1;
+    if (bDate > aDate) return 1;
+    return 0;
+  });
+
+  // set the y-values in each subtree
+  // such that oldest subtrees are at the top
+  let currentMaxY = 0;
+  while (subtreeStack.length) {
+    const nextSubtree = subtreeStack.shift();
+    nextSubtree.subtreeNodes.forEach(node => node.n.yvalue = currentMaxY++)
+  }
+}
+
 
 export const formatDivergence = (divergence) => {
   return divergence > 1 ?

src/components/tree/phyloTree/layouts.js

@@ -34,9 +34,8 @@ export const setLayout = function setLayout(layout) {
   timerEnd("setLayout");
 };
 
-
 /**
- * assignes x,y coordinates for a rectancular layout
+ * assigns x,y coordinates for a rectangular layout
  * @return {null}
  */
 export const rectangularLayout = function rectangularLayout() {

src/components/tree/phyloTree/renderers.js

@@ -172,7 +172,7 @@ export const drawBranches = function drawBranches() {
   if (!("branchTee" in this.groups)) {
     this.groups.branchTee = this.svg.append("g").attr("id", "branchTee");
   }
-  if (this.layout === "clock" || this.layout === "unrooted") {
+  if (this.layout === "clock" || this.layout === "unrooted" || params.isTreeSplitByTrait) {
     this.groups.branchTee.selectAll("*").remove();
   } else {
     this.groups.branchTee

src/components/tree/reactD3Interface/change.js

@@ -92,6 +92,18 @@ export const changePhyloTreeViaPropsComparison = (mainTree, phylotree, oldProps,
     args.svgHasChangedDimensions = true;
   }
 
+  /* split tree by colored-by trait */
+  /* just copies into tree; will be null/empty string/not defined if no trait */
+  if (newProps.tree.splitTreeByTrait !== oldProps.tree.splitTreeByTrait)
+  {
+    args.splitTreeByTrait = newProps.tree.splitTreeByTrait;
+    if (args.splitTreeByTrait == null) {// if the split is being unset, reset the layout
+    // todo: could use updateLayout instead?
+      args.newLayout = newProps.layout;
+      args.resetTreeYValues = true;
+    }
+  }
+
   const change = Object.keys(args).length;
   if (change) {
     args.animationInProgress = newProps.animationPlayPauseButton === "Pause";

src/reducers/controls.js

@@ -81,7 +81,8 @@ export const getDefaultControlsState = () => {
     sidebarOpen: initialSidebarState.sidebarOpen,
     treeLegendOpen: undefined,
     mapLegendOpen: undefined,
-    showOnlyPanels: false
+    showOnlyPanels: false,
+    splitTreeByTrait: undefined
   };
 };
 

src/reducers/tree.js

@@ -24,7 +24,8 @@ export const getDefaultTreeState = () => {
     totalStateCounts: {},
     availableBranchLabels: [],
     selectedStrain: undefined,
-    selectedClade: undefined
+    selectedClade: undefined,
+    splitTreeByTrait: null
   };
 };
 
@@ -75,6 +76,13 @@ const Tree = (state = getDefaultTreeState(), action) => {
         }
       });
       return state;
+    case types.TOGGLE_SPLIT_TREE:
+      // todo: what if nothing is colored-by? can that happen?
+      // note that there's no reason technically not to split by a state that isn't the colored-by
+      // state, it just doesn't make as much UI sense  
+      return Object.assign({}, state, {
+        splitTreeByTrait: action.splitTreeByTrait
+      });
     default:
       return state;
   }