Improve second tree loading & parsing

CHANGELOG.md

@@ -1,9 +1,12 @@
 # Changelog
 
-* remove missing warning then dataset JSON has no `.meta.geo_resolutions` field ([#1791](https://github.com/nextstrain/auspice/pull/1791))
+* A number of improvements when viewing multiple trees (tangletrees) ([#1788](https://github.com/nextstrain/auspice/pull/1788))
+  * Attributes only present on the RHS tree are now available as filter options, as well as genotypes and node names unique to the RHS tree.
+  * Filter badges (shown in the header) now indicate how many matches are present in both trees (formerly only the LHS tree was considered)
+  * Branch labels unique to the RHS tree are now available
+* remove missing warning when dataset JSON has no `.meta.geo_resolutions` field ([#1791](https://github.com/nextstrain/auspice/pull/1791))
 * Add support for Node.js version 22. ([#1779](https://github.com/nextstrain/auspice/pull/1779))
 
-
 ## version 2.54.3 - 2024/06/12
 
 

src/actions/recomputeReduxState.js

@@ -690,26 +690,6 @@ const modifyTreeStateVisAndBranchThickness = (oldState, zoomSelected, controlsSt
   return newState;
 };
 
-const removePanelIfPossible = (panels, name) => {
-  const idx = panels.indexOf(name);
-  if (idx !== -1) {
-    panels.splice(idx, 1);
-  }
-};
-
-const modifyControlsViaTreeToo = (controls, name) => {
-  controls.showTreeToo = name;
-  controls.showTangle = true;
-  controls.layout = "rect"; /* must be rectangular for two trees */
-  controls.panelsToDisplay = controls.panelsToDisplay.slice();
-  removePanelIfPossible(controls.panelsToDisplay, "map");
-  removePanelIfPossible(controls.panelsToDisplay, "entropy");
-  removePanelIfPossible(controls.panelsToDisplay, "frequencies");
-  controls.canTogglePanelLayout = false;
-  controls.panelLayout = "full";
-  return controls;
-};
-
 /**
  * The v2 JSON spec defines colorings as a list, so that order is guaranteed.
  * Prior to this, we used a dict, where key insertion order is (guaranteed? essentially always?)
@@ -877,14 +857,7 @@ export const createStateFromQueryOrJSONs = ({
     tree.name = mainTreeName;
     metadata.mainTreeNumTips = calcTotalTipsInTree(tree.nodes);
     if (secondTreeDataset) {
-      treeToo = treeJsonToState(secondTreeDataset.tree);
-      castIncorrectTypes(metadata, treeToo);
-      treeToo.debug = "RIGHT";
-      treeToo.name = secondTreeName;
-      updateMetadataStateViaSecondTree(metadata, secondTreeDataset, entropy?.genomeMap)
-
-      /* TODO: calc & display num tips in 2nd tree */
-      // metadata.secondTreeNumTips = calcTotalTipsInTree(treeToo.nodes);
+      ({treeToo, metadata} = instantiateSecondTree(secondTreeDataset, metadata, entropy?.genomeMap, secondTreeName));
     }
 
     /* new controls state - don't apply query yet (or error check!) */
@@ -965,11 +938,7 @@ export const createStateFromQueryOrJSONs = ({
   tree = modifyTreeStateVisAndBranchThickness(tree, query.label, controls, dispatch);
 
   if (treeToo && treeToo.loaded) {
-    treeToo.nodeColorsVersion = tree.nodeColorsVersion;
-    treeToo.nodeColors = calcNodeColor(treeToo, controls.colorScale);
-    treeToo = modifyTreeStateVisAndBranchThickness(treeToo, undefined, controls, dispatch);
-    controls = modifyControlsViaTreeToo(controls, treeToo.name);
-    treeToo.tangleTipLookup = constructVisibleTipLookupBetweenTrees(tree.nodes, treeToo.nodes, tree.visibility, treeToo.visibility);
+    treeToo = updateSecondTree(tree, treeToo, controls, dispatch)
   }
 
   /* we can only calculate which legend items we wish to display _after_ the visibility has been calculated */
@@ -1044,33 +1013,64 @@ export const createTreeTooState = ({
   /* TODO: reconcile choices (filters, colorBys etc) with this new tree */
   /* TODO: reconcile query with visibility etc */
 
-  const metadata = {...oldState.metadata};
-  updateMetadataStateViaSecondTree(metadata, json, oldState.entropy?.genomeMap);
-
-  let controls = {...oldState.controls};
   const tree = Object.assign({}, oldState.tree);
   tree.name = originalTreeUrl;
-  let treeToo = treeJsonToState(json.tree);
-  treeToo.name = secondTreeUrl;
-  treeToo.debug = "RIGHT";
-  controls = modifyControlsStateViaTree(controls, tree, treeToo, oldState.metadata.colorings);
-  controls = modifyControlsViaTreeToo(controls, secondTreeUrl);
-  treeToo = modifyTreeStateVisAndBranchThickness(treeToo, undefined, controls, dispatch);
+  let {treeToo, metadata} = instantiateSecondTree(json, oldState.metadata, oldState.entropy?.genomeMap, secondTreeUrl)
 
-  /* calculate colours if loading from JSONs or if the query demands change */
-  const colorScale = calcColorScale(controls.colorBy, controls, tree, treeToo, metadata);
-  const nodeColors = calcNodeColor(treeToo, colorScale);
-  tree.nodeColors = calcNodeColor(tree, colorScale); // also update main tree's colours
-  tree.nodeColorsVersion++;
-
-  controls.colorScale = colorScale;
+  /* recompute the controls state now that we have new data */
+  const controls = modifyControlsStateViaTree({...oldState.controls}, tree, treeToo, oldState.metadata.colorings);
+  /* recalculate the color scale with updated tree data */
+  controls.colorScale = calcColorScale(controls.colorBy, controls, tree, treeToo, metadata);
   controls.colorByConfidence = doesColorByHaveConfidence(controls, controls.colorBy);
-  treeToo.nodeColorsVersion = colorScale.version;
-  treeToo.nodeColors = nodeColors;
+  /* and update the color scale as applied to the LHS tree */
+  tree.nodeColors = calcNodeColor(tree, controls.colorScale);
+  tree.nodeColorsVersion++;
 
-  treeToo.tangleTipLookup = constructVisibleTipLookupBetweenTrees(
-    tree.nodes, treeToo.nodes, tree.visibility, treeToo.visibility
-  );
+  treeToo = updateSecondTree(tree, treeToo, controls, dispatch);
 
   return {tree, treeToo, controls, metadata};
 };
+
+
+/**
+ * Code which is common to both loading a second tree within `createStateFromQueryOrJSONs` and
+ * `createTreeTooState`.
+ */
+function instantiateSecondTree(secondTreeDataset, metadata, genomeMap, secondTreeName) {
+  const treeToo = treeJsonToState(secondTreeDataset.tree);
+  castIncorrectTypes(metadata, treeToo);
+  treeToo.debug = "RIGHT";
+  treeToo.name = secondTreeName;
+  updateMetadataStateViaSecondTree({...metadata}, secondTreeDataset, genomeMap)
+
+  const secondTreeColorings = convertColoringsListToDict(secondTreeDataset.meta?.colorings || []);
+  const stateCountAttrs = gatherTraitNames(treeToo.nodes, secondTreeColorings);
+  treeToo.totalStateCounts = countTraitsAcrossTree(treeToo.nodes, stateCountAttrs, false, true);
+
+  /* TODO: calc & display num tips in 2nd tree */
+  // metadata.secondTreeNumTips = calcTotalTipsInTree(treeToo.nodes);
+
+  return {treeToo, metadata}
+}
+
+/**
+ * Update colours and control state options. This function requires that the controls state
+ * has been instantiated (e.g. the colorScale has been computed)
+ */
+function updateSecondTree(tree, treeToo, controls, dispatch) {
+  treeToo.nodeColorsVersion = tree.nodeColorsVersion;
+  treeToo.nodeColors = calcNodeColor(treeToo, controls.colorScale);
+  treeToo = modifyTreeStateVisAndBranchThickness(treeToo, undefined, controls, dispatch);
+  treeToo.tangleTipLookup = constructVisibleTipLookupBetweenTrees(tree.nodes, treeToo.nodes, tree.visibility, treeToo.visibility);
+
+  /* modify controls */
+  controls.showTreeToo = treeToo.name;
+  controls.showTangle = true;
+  controls.layout = "rect"; /* must be rectangular for two trees */
+  controls.panelsToDisplay = controls.panelsToDisplay
+    .filter((name) => !["map", "entropy", "frequencies"].includes(name));
+  controls.canTogglePanelLayout = false;
+  controls.panelLayout = "full";
+
+  return treeToo;
+}

src/components/controls/choose-branch-labelling.js

@@ -11,7 +11,10 @@ import Toggle from "./toggle";
 @connect((state) => ({
   selected: state.controls.selectedBranchLabel,
   showAll: state.controls.showAllBranchLabels,
-  available: state.tree.availableBranchLabels,
+  available: Array.from(
+    (new Set(state.tree.availableBranchLabels))
+      .union(new Set(state.treeToo?.availableBranchLabels ?? []))
+  ),
   canRenderBranchLabels: state.controls.canRenderBranchLabels
 }))
 class ChooseBranchLabelling extends React.Component {

src/components/controls/filter.js

@@ -26,6 +26,8 @@ const DEBOUNCE_TIME = 200;
     totalStateCounts: state.tree.totalStateCounts,
     canFilterByGenotype: !!state.entropy.genomeMap,
     nodes: state.tree.nodes,
+    nodesSecondTree: state.treeToo?.nodes,
+    totalStateCountsSecondTree: state.treeToo?.totalStateCounts,
     measurementsFieldsMap: state.measurements.collectionToDisplay.fields,
     measurementsFiltersMap: state.measurements.collectionToDisplay.filters,
     measurementsFilters: state.controls.measurementsFilters
@@ -66,16 +68,24 @@ class FilterData extends React.Component {
      * colorings). Within each trait, the values are alphabetical
      */
     const coloringKeys = Object.keys(this.props.colorings||{});
-    const unorderedTraitNames = Object.keys(this.props.totalStateCounts);
+    const unorderedTraitNames = [
+      ...Object.keys(this.props.totalStateCounts),
+      ...Object.keys(this.props.totalStateCountsSecondTree),
+    ]
     const traitNames = [
       ...coloringKeys.filter((name) => unorderedTraitNames.includes(name)),
       ...unorderedTraitNames.filter((name) => !coloringKeys.includes(name))
     ]
     for (const traitName of traitNames) {
-      const traitData = this.props.totalStateCounts[traitName];
+      const traitData = new Set([
+        ...(this.props.totalStateCounts[traitName]?.keys() || []),
+        ...(this.props.totalStateCountsSecondTree?.[traitName]?.keys() || []),
+      ]);
+
+      this.props.totalStateCounts[traitName];
       const traitTitle = this.getFilterTitle(traitName);
       const filterValuesCurrentlyActive = new Set((this.props.activeFilters[traitName] || []).filter((x) => x.active).map((x) => x.value));
-      for (const traitValue of Array.from(traitData.keys()).sort()) {
+      for (const traitValue of Array.from(traitData).sort()) {
         if (filterValuesCurrentlyActive.has(traitValue)) continue;
         options.push({
           label: `${traitTitle} → ${traitValue}`,
@@ -89,7 +99,12 @@ class FilterData extends React.Component {
      * mutations
      */
     if (this.props.canFilterByGenotype) {
-      Array.from(collectGenotypeStates(this.props.nodes))
+      const observedGenotypes = collectGenotypeStates(this.props.nodes); // set of "nuc:123A", "S:418K", etc
+      const observedGenotypesSecondTree = this.props.nodesSecondTree ?
+        collectGenotypeStates(this.props.nodesSecondTree).difference(observedGenotypes) :
+        new Set();
+      Array.from(observedGenotypes)
+        .concat(Array.from(observedGenotypesSecondTree))
         .sort()
         .forEach((o) => {
           options.push({
@@ -99,14 +114,29 @@ class FilterData extends React.Component {
         });
     }
 
-    this.props.nodes
+    /**
+     * Add all (terminal) node names, calling each a "sample"
+     */
+    const sampleNames = this.props.nodes
       .filter((n) => !n.hasChildren)
-      .forEach((n) => {
+      .map((n) => n.name);
+    sampleNames.forEach((name) => {
         options.push({
-          label: `sample → ${n.name}`,
-          value: [strainSymbol, n.name]
+          label: `sample → ${name}`,
+          value: [strainSymbol, name]
         });
       });
+    if (this.props.nodesSecondTree) {
+      const seenNames = new Set(sampleNames);
+      this.props.nodesSecondTree
+        .filter((n) => !n.hasChildren && !seenNames.has(n.name))
+        .forEach((n) => {
+          options.push({
+            label: `sample → ${n.name}`,
+            value: [strainSymbol, n.name]
+          });
+        });
+    }
 
     if (this.props.measurementsOn && this.props.measurementsFiltersMap && this.props.measurementsFieldsMap) {
       this.props.measurementsFiltersMap.forEach(({values}, filterField) => {

src/components/info/filtersSummary.js

@@ -33,6 +33,7 @@ const closeBracketSmall = <span style={{fontSize: "1.8rem", fontWeight: 300, pad
     metadata: state.metadata,
     nodes: state.tree.nodes,
     totalStateCounts: state.tree.totalStateCounts,
+    totalStateCountsSecondTree: state.treeToo?.totalStateCounts,
     visibility: state.tree.visibility,
     selectedClade: state.tree.selectedClade,
     dateMin: state.controls.dateMin,
@@ -86,7 +87,19 @@ class FiltersSummary extends React.Component {
       .sort((a, b) => a.value < b.value ? -1 : a.value > b.value ? 1 : 0)
       .map((item) => {
         let label = `${item.value}`;
-        if (filterName!==strainSymbol) label+= ` (${this.props.totalStateCounts[filterName].get(item.value)})`;
+        if (filterName!==strainSymbol) {
+          /* Add the _total_ occurrences in parentheses. We don't compute the intersections / visible values here -
+          e.g. we can have "England (5)" "North America (100)" even though such an intersection will deselect everything.
+          If we have two trees shown we show both values.
+          */
+          const tree1count = this.props.totalStateCounts[filterName]?.get(item.value) ?? 0;
+          if (this.props.totalStateCountsSecondTree) {
+            const tree2count = this.props.totalStateCountsSecondTree[filterName]?.get(item.value) ?? 0;
+            label+=` (L: ${tree1count}, R: ${tree2count})`;
+          } else {
+            label+=` (${tree1count})`;
+          }
+        }
         return this.createIndividualBadge({filterName, item, label, onHoverMessage});
       });
   }