Streamtrees

package.json

@@ -43,6 +43,7 @@
     "@babel/preset-typescript": "^7.21.0",
     "@hot-loader/react-dom": "^16.13.0",
     "@reduxjs/toolkit": "^1.9.7",
+    "@stdlib/stats-base-dists-normal-pdf": "^0.2.2",
     "argparse": "^1.0.10",
     "babel-loader": "^8.0.4",
     "babel-plugin-lodash": "^3.3.4",

src/actions/colors.js

@@ -4,6 +4,7 @@ import { calcColorScale } from "../util/colorScale";
 import { timerStart, timerEnd } from "../util/perf";
 import { changeEntropyCdsSelection } from "./entropy";
 import { updateFrequencyDataDebounced } from "./frequencies";
+import { processStreams } from "../util/treeStreams";
 import * as types from "./types";
 
 /* providedColorBy: undefined | string */
@@ -27,6 +28,10 @@ export const changeColorBy = (providedColorBy = undefined) => {
 
     dispatch(changeEntropyCdsSelection(colorBy));
 
+    if (Object.keys(tree.streams).length) {
+      processStreams(tree.streams, tree.nodes, tree.visibility, controls.distanceMeasure, colorScale, {skipPivots: true})
+    }
+
     dispatch({
       type: types.NEW_COLORS,
       colorBy,

src/actions/layout.js

@@ -1,5 +1,6 @@
 import { CHANGE_LAYOUT } from "./types";
 import { validateScatterVariables, addScatterAxisInfo} from "../util/scatterplotHelpers";
+import { toggleStreamTree } from "./treeStreams";
 
 /**
  * Redux Thunk to change a layout, including aspects of the scatterplot / clock layouts.
@@ -9,6 +10,11 @@ export const changeLayout = ({layout, showBranches, showRegression, x, xLabel, y
     if (window.NEXTSTRAIN && window.NEXTSTRAIN.animationTickReference) return;
     const { controls, tree, metadata } = getState();
 
+    if (controls.showStreamTrees) {
+      // Note that these multiple dispatches prevent the nice d3 layout change
+      dispatch(toggleStreamTree());
+    }
+
     if (layout==="rect" || layout==="unrooted" || layout==="radial") {
       dispatch({type: CHANGE_LAYOUT, layout, scatterVariables: controls.scatterVariables, canRenderBranchLabels: true});
       return;

src/actions/recomputeReduxState.js

@@ -23,6 +23,7 @@ import { collectAvailableTipLabelOptions } from "../components/controls/choose-t
 import { hasMultipleGridPanels } from "./panelDisplay";
 import { strainSymbolUrlString } from "../middleware/changeURL";
 import { combineMeasurementsControlsAndQuery, encodeMeasurementColorBy, loadMeasurements } from "./measurements";
+import { processStreams, labelStreamMembership } from "../util/treeStreams";
 
 export const doesColorByHaveConfidence = (controlsState, colorBy) =>
   controlsState.coloringsPresentOnTreeWithConfidence.has(colorBy);
@@ -460,6 +461,19 @@ const modifyControlsStateViaTree = (state, tree, treeToo, colorings) => {
     state.selectedBranchLabel = "clade";
   }
 
+  const candidateStreamBranchLabels = [
+    "stream_label",
+    "stream",
+    // "clade",
+  ];
+  for (const key of candidateStreamBranchLabels) {
+    if (tree.availableBranchLabels.includes(key)) {
+      state.streamTreeBranchLabel = key;
+      state.showStreamTrees = true; // TODO XXX - remove - here for dev purposes only
+    }
+  }
+
+
   state.temporalConfidence = {exists: num_date_confidence, display: num_date_confidence, on: false};
 
   return state;
@@ -1034,6 +1048,17 @@ export const createStateFromQueryOrJSONs = ({
   /* if query.label is undefined then we intend to zoom to the root */
   tree = modifyTreeStateVisAndBranchThickness(tree, query.label, controls, dispatch);
 
+  /* scan the tree and identify / label streams for the currently chosen
+  Note this happens irrespective of `showStreamTrees`, but we could defer this
+  to improve first load time */
+  if (controls.streamTreeBranchLabel!=='none') {
+    tree.streams = labelStreamMembership(tree.nodes[0], controls.streamTreeBranchLabel)
+  }
+
+  if (controls.showStreamTrees && !!Object.keys(tree.streams).length) {
+    processStreams(tree.streams, tree.nodes, tree.visibility, controls.distanceMeasure, controls.colorScale)
+  }
+
   if (treeToo && treeToo.loaded) {
     treeToo = updateSecondTree(tree, treeToo, controls, dispatch)
   }

src/actions/tree.ts

@@ -15,6 +15,7 @@ import { PhyloNode } from "../components/tree/phyloTree/types";
 import { Metadata } from "../metadata";
 import { ThunkFunction } from "../store";
 import { ReduxNode, TreeState } from "../reducers/tree/types";
+import { processStreams } from "../util/treeStreams";
 
 type RootIndex = number | undefined
 
@@ -135,6 +136,11 @@ export const updateVisibleTipsAndBranchThicknesses = ({
       /* tip selected is the same as the first tree - the reducer uses that */
     }
 
+    if (Object.keys(tree.streams).length) {
+      // recomputes them even if they're toggled off
+      processStreams(tree.streams, tree.nodes, dispatchObj.visibility,  controls.distanceMeasure, controls.colorScale, {skipPivots: true, skipCategories: true});
+    }
+
     /* Changes in visibility require a recomputation of which legend items we wish to display */
     dispatchObj.visibleLegendValues = createVisibleLegendValues({
       colorBy: controls.colorBy,
@@ -213,6 +219,11 @@ export const changeDateFilter = ({
       visibilityToo: dispatchObj.visibilityToo
     });
 
+    if (Object.keys(tree.streams).length) {
+      // recomputes them even if they're toggled off
+      processStreams(tree.streams, tree.nodes, dispatchObj.visibility,  controls.distanceMeasure, controls.colorScale, {skipPivots: true, skipCategories: true});
+    }
+
     /* D I S P A T C H */
     dispatch(dispatchObj);
     updateEntropyVisibility(dispatch, getState);
@@ -474,3 +485,15 @@ export const explodeTree = (
     });
   };
 };
+
+
+
+export function changeDistanceMeasure(metric: "num_date"|"div"): ThunkFunction {
+  return function(dispatch, getState) {
+    const {controls, tree} = getState();
+    if (Object.keys(tree.streams).length) {
+      processStreams(tree.streams, tree.nodes, tree.visibility,  metric, controls.colorScale, {skipCategories: true})
+    }
+    dispatch({type: types.CHANGE_DISTANCE_MEASURE, data: metric})
+  }
+}

src/actions/treeStreams.ts

@@ -0,0 +1,67 @@
+import { TOGGLE_STREAM_TREE, CHANGE_STREAM_TREE_BRANCH_LABEL } from "./types";
+import { processStreams, labelStreamMembership, isNodeWithinAnotherStream } from "../util/treeStreams";
+import { availableStreamTreeBranchLabels } from "../components/controls/choose-stream-trees";
+import { getParentStream } from "../components/tree/phyloTree/helpers";
+import { updateVisibleTipsAndBranchThicknesses } from "./tree";
+import { warningNotification } from "./notifications";
+
+export function toggleStreamTree() {
+  return function(dispatch, getState) {
+    const {controls, tree} = getState();
+    const showStreamTrees = !controls.showStreamTrees; // new state
+
+    if (showStreamTrees===false) {
+      dispatch({type: TOGGLE_STREAM_TREE, showStreamTrees});
+      return;
+    }
+
+    if (controls.streamTreeBranchLabel==='none') {
+      // toggle switched on without an already set branch label
+      dispatch(changeStreamTreeBranchLabel(availableStreamTreeBranchLabels(tree.availableBranchLabels)[0]));
+      return;
+    }
+
+    const currentRoot = tree.nodes[tree.idxOfInViewRootNode]
+    if (currentRoot && currentRoot.inStream && !currentRoot.streamName) {
+      /**
+       * This block indicates a situation where we aren't viewing streams, but the current in-view root is _within_ a stream and
+       * we're asking to show streams. We can't display a stream starting half-way through the stream, so we can either
+       * (a) prevent this action from happening or (b) zoom back out to show the entire stream we're currently within.
+       * We do (b) and use a `setTimeout` to ensure that we should the zooming out behaviour before the stream is toggled on.
+       * Timeout's are prone to bugs and so improvements here are welcome. TODO XXX
+       */
+      const newRootNode = getParentStream(currentRoot);
+      dispatch(updateVisibleTipsAndBranchThicknesses({root: [newRootNode.arrayIdx, undefined]}));
+      window.setTimeout(() => dispatch(toggleStreamTree()), 500)
+      return;
+    }
+
+    processStreams(tree.streams, tree.nodes, tree.visibility, controls.distanceMeasure, controls.colorScale)
+    dispatch({type: TOGGLE_STREAM_TREE, showStreamTrees})
+  }
+} 
+
+
+export function changeStreamTreeBranchLabel(newLabel) {
+  return function(dispatch, getState) {
+    const {controls, tree} = getState();
+
+    if (isNodeWithinAnotherStream(tree.nodes[tree.idxOfInViewRootNode], newLabel)) {
+      dispatch(warningNotification({message: `Cannot switch streams to ${newLabel} as the subtree we're viewing would be inside a stream`}));
+      return;
+    }
+
+    const streams = labelStreamMembership(tree.nodes[0], newLabel);
+
+    const showStreamTrees = newLabel!=='none';
+    if (showStreamTrees && !!Object.keys(streams).length) {
+      processStreams(streams, tree.nodes, tree.visibility, controls.distanceMeasure, controls.colorScale);
+    }
+    dispatch({
+      type: CHANGE_STREAM_TREE_BRANCH_LABEL,
+      streams,
+      showStreamTrees,
+      streamTreeBranchLabel: newLabel
+    })
+  }
+}

src/actions/types.js

@@ -64,3 +64,5 @@ export const TOGGLE_SHOW_ALL_BRANCH_LABELS = "TOGGLE_SHOW_ALL_BRANCH_LABELS";
 export const TOGGLE_MOBILE_DISPLAY = "TOGGLE_MOBILE_DISPLAY";
 export const SELECT_NODE = "SELECT_NODE";
 export const DESELECT_NODE = "DESELECT_NODE";
+export const TOGGLE_STREAM_TREE = "TOGGLE_STREAM_TREE";
+export const CHANGE_STREAM_TREE_BRANCH_LABEL = "CHANGE_STREAM_TREE_BRANCH_LABEL";

src/components/controls/choose-explode-attr.js

@@ -16,7 +16,8 @@ import CustomSelect from "./customSelect";
   selected: state.controls.explodeAttr,
   available: state.metadata.colorings,
   showThisUI: !state.controls.showTreeToo,
-  mobileDisplay: state.general.mobileDisplay
+  mobileDisplay: state.general.mobileDisplay,
+  showStreamTrees: state.controls.showStreamTrees,
 }))
 class ChooseExplodeAttr extends React.Component {
   constructor(props) {
@@ -35,6 +36,10 @@ class ChooseExplodeAttr extends React.Component {
   }
   render() {
     if (!this.props.showThisUI) return null;
+    // Don't display the sidebar UI if we're showing stream trees
+    // (we could follow the UI example of the stream tree toggle & indicate _why_)
+    if (this.props.showStreamTrees) return null;
+
     const { t, tooltip } = this.props;
     const selectOptions = this.gatherAttrs();
     return (

src/components/controls/choose-metric.js

@@ -1,9 +1,8 @@
 import React from "react";
 import { connect } from "react-redux";
 import { withTranslation } from "react-i18next";
-import { CHANGE_DISTANCE_MEASURE } from "../../actions/types";
 import { analyticsControlsEvent } from "../../util/googleAnalytics";
-import { toggleTemporalConfidence } from "../../actions/tree";
+import { toggleTemporalConfidence, changeDistanceMeasure } from "../../actions/tree";
 import { SidebarSubtitle, SidebarButton } from "./styles";
 import Toggle from "./toggle";
 
@@ -34,7 +33,7 @@ class ChooseMetric extends React.Component {
           selected={this.props.distanceMeasure === "num_date"}
           onClick={() => {
             analyticsControlsEvent("tree-metric-temporal");
-            this.props.dispatch({ type: CHANGE_DISTANCE_MEASURE, data: "num_date" });
+            this.props.dispatch(changeDistanceMeasure("num_date"));
           }}
         >
           {t("sidebar:time")}
@@ -44,7 +43,7 @@ class ChooseMetric extends React.Component {
           selected={this.props.distanceMeasure === "div"}
           onClick={() => {
             analyticsControlsEvent("tree-metric-temporal");
-            this.props.dispatch({ type: CHANGE_DISTANCE_MEASURE, data: "div" });
+            this.props.dispatch(changeDistanceMeasure("div"));
           }}
         >
           {t("sidebar:divergence")}

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

@@ -11,6 +11,7 @@ import CustomSelect from "./customSelect";
 @connect((state) => {
   return {
     available: state.controls.available,
+    showStreamTrees: state.controls.showStreamTrees,
     treeName: state.tree.name,
     showTreeToo: state.controls.showTreeToo /* this is the name of the second tree if one is selected */
   };
@@ -35,6 +36,9 @@ class ChooseSecondTree extends React.Component {
 
     // Don't display the sidebar UI if we're just going to display an empty dropdown!
     if (!options.length) return null;
+    // Don't display the sidebar UI if we're showing stream trees
+    // (we could follow the UI example of the stream tree toggle & indicate _why_)
+    if (this.props.showStreamTrees) return null;
 
     if (this.props.showTreeToo) options.unshift("REMOVE");