Merge pull request #1785 from nextstrain/james/tangletree-genotype-co…

…louring

Tangletree genotype colouring

CHANGELOG.md

@@ -1,7 +1,9 @@
 # Changelog
 
-## version 2.54.1 - 2024/06/10
+* Tangletrees can now be coloured by genotype (previously such a colouring would only work for the LHS tree). This requires the genome annotation (in the dataset JSON) to be identical across both datasets. This can be especially useful when comparing trees generated from the same sequences or a similar set of sequences in order to understand the differences in tree structure. ([#1785](https://github.com/nextstrain/auspice/pull/1783))
+* Bugfix: The legend entries shown for a tangletree may not have shown values only observed in the RHS tree when the dataset was first loaded. ([#1785](https://github.com/nextstrain/auspice/pull/1783))
 
+## version 2.54.1 - 2024/06/10
 
 * Fixed a big bug where clicking on tips (and shift-clicking on branches) on the RHS tree in a tanglegram would bring up a modal detailing a node in the LHS tree. ([#1783](https://github.com/nextstrain/auspice/pull/1783))
 * Fixed a small bug where branch labels prevented you from hovering on the branch itself, a situation that was more common in tangletrees. ([#1783](https://github.com/nextstrain/auspice/pull/1783))

src/actions/loadData.js

@@ -50,7 +50,7 @@ export const loadSecondTree = (secondTreeUrl, firstTreeUrl) => async (dispatch,
     dispatch(explodeTree(undefined));
   }
 
-  const newState = createTreeTooState({treeTooJSON: secondJson.tree, oldState, originalTreeUrl: firstTreeUrl, secondTreeUrl: secondTreeUrl, dispatch});
+  const newState = createTreeTooState({json: secondJson, oldState, originalTreeUrl: firstTreeUrl, secondTreeUrl: secondTreeUrl, dispatch});
   dispatch({type: types.TREE_TOO_DATA, ...newState});
 };
 

src/actions/recomputeReduxState.js

@@ -1,5 +1,5 @@
 import queryString from "query-string";
-import { cloneDeep } from 'lodash';
+import { cloneDeep, isEqualWith } from 'lodash';
 import { numericToCalendar, calendarToNumeric } from "../util/dateHelpers";
 import { reallySmallNumber, twoColumnBreakpoint, defaultColorBy, defaultGeoResolution, defaultDateRange, nucleotide_gene, strainSymbol, genotypeSymbol } from "../util/globals";
 import { calcBrowserDimensionsInitialState } from "../reducers/browserDimensions";
@@ -12,7 +12,7 @@ import { countTraitsAcrossTree, calcTotalTipsInTree, gatherTraitNames } from "..
 import { calcEntropyInView } from "../util/entropy";
 import { treeJsonToState } from "../util/treeJsonProcessing";
 import { castIncorrectTypes } from "../util/castJsonTypes";
-import { entropyCreateState } from "../util/entropyCreateStateFromJsons";
+import { entropyCreateState, genomeMap as createGenomeMap } from "../util/entropyCreateStateFromJsons";
 import { calcNodeColor } from "../util/colorHelpers";
 import { calcColorScale, createVisibleLegendValues } from "../util/colorScale";
 import { computeMatrixFromRawData, checkIfNormalizableFromRawData } from "../util/processFrequencies";
@@ -800,6 +800,47 @@ const createMetadataStateFromJSON = (json) => {
   return metadata;
 };
 
+/**
+ * Conceptually similar to `createMetadataStateFromJSON` but here looking at metadata from
+ * the (optional) second tree and only considers certain properties at the moment.
+ */
+function updateMetadataStateViaSecondTree(metadata, json, genomeMap) {
+
+  // For genotype colourings across multiple trees we need to know if the genome
+  // maps are identical¹ across both trees
+  //
+  // ¹ This could be relaxed in the future - currently we enforce that the order
+  //   of genes matches. We could also make this more fine grained and allow the
+  //   2nd tree to have a subset of CDSs (wrt the main tree), with genotypes
+  //   only working for the shared CDSs.     
+  if (genomeMap && json.meta.genome_annotations) {
+    try {
+      metadata.identicalGenomeMapAcrossBothTrees = isEqualWith(
+        genomeMap,
+        createGenomeMap(json.meta.genome_annotations),
+        (objValue, othValue, indexOrKey) => {
+          if (indexOrKey==='color') return true; // don't compare CDS colors!
+          // don't compare metadata section as there may be Infinities here
+          // (and if everything else is equal then the metadata will be the same too)
+          if (indexOrKey==='metadata') return true;
+          return undefined; // use lodash's default comparison
+        }
+      )
+    } catch (e) {
+      if (e instanceof Error) console.error(e.message);
+    }
+    if (!metadata.identicalGenomeMapAcrossBothTrees) {
+      console.warn("Heads up! The two trees have different genome_annotations and thus genotype colorings will only be applied to the LHS tree")
+    }
+  }
+
+  if (metadata.identicalGenomeMapAcrossBothTrees && json.root_sequence) {
+    /* in-line root sequence (not sidecar) */
+    metadata.rootSequenceSecondTree = json.root_sequence;
+  }
+
+}
+
 export const getNarrativePageFromQuery = (query, narrative) => {
   let n = parseInt(query.n, 10) || 0;
   /* If the query has defined a block which doesn't exist then default to n=0 */
@@ -842,6 +883,8 @@ export const createStateFromQueryOrJSONs = ({
       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);
     }
@@ -940,7 +983,7 @@ export const createStateFromQueryOrJSONs = ({
     treeNodes: tree.nodes,
     treeTooNodes: treeToo ? treeToo.nodes : undefined,
     visibility: tree.visibility,
-    visibilityToo: treeToo ? treeToo.visibilityToo : undefined
+    visibilityToo: treeToo?.visibility,
   });
 
   /* calculate entropy in view */
@@ -994,26 +1037,30 @@ export const createStateFromQueryOrJSONs = ({
 };
 
 export const createTreeTooState = ({
-  treeTooJSON, /* raw json data */
+  json, /* raw json data */
   oldState,
   originalTreeUrl,
   secondTreeUrl, /* treeToo URL */
   dispatch
 }) => {
   /* TODO: reconcile choices (filters, colorBys etc) with this new tree */
   /* TODO: reconcile query with visibility etc */
-  let controls = oldState.controls;
+
+  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(treeTooJSON);
+  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);
 
   /* calculate colours if loading from JSONs or if the query demands change */
-  const colorScale = calcColorScale(controls.colorBy, controls, tree, treeToo, oldState.metadata);
+  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++;
@@ -1027,5 +1074,5 @@ export const createTreeTooState = ({
     tree.nodes, treeToo.nodes, tree.visibility, treeToo.visibility
   );
 
-  return {tree, treeToo, controls};
+  return {tree, treeToo, controls, metadata};
 };

src/reducers/metadata.js

@@ -10,6 +10,8 @@ const Metadata = (state = {
   loaded: false, /* see comment in the sequences reducer for explanation */
   metadata: null,
   rootSequence: undefined,
+  identicalGenomeMapAcrossBothTrees: false,
+  rootSequenceSecondTree: undefined,
   colorOptions // this can't be removed as the colorScale currently runs before it should
 }, action) => {
   switch (action.type) {
@@ -18,6 +20,7 @@ const Metadata = (state = {
         loaded: false
       });
     case types.URL_QUERY_CHANGE_WITH_COMPUTED_STATE:
+    case types.TREE_TOO_DATA:
     case types.CLEAN_START:
       return action.metadata;
     case types.ADD_EXTRA_METADATA: {
@@ -41,6 +44,11 @@ const Metadata = (state = {
     }
     case types.SET_ROOT_SEQUENCE:
       return {...state, rootSequence: action.data};
+    case types.REMOVE_TREE_TOO:
+      return Object.assign({}, state, {
+        identicalGenomeMapAcrossBothTrees: false,
+        rootSequenceSecondTree: undefined,
+      });
     default:
       return state;
   }

src/util/colorScale.js

@@ -39,10 +39,13 @@ export const calcColorScale = (colorBy, controls, tree, treeToo, metadata) => {
     if (isColorByGenotype(colorBy)) {
       genotype = decodeColorByGenotype(colorBy);
       setGenotype(tree.nodes, genotype.gene, genotype.positions, metadata.rootSequence); /* modifies nodes recursively */
+      if (treeToo && metadata.identicalGenomeMapAcrossBothTrees) {
+        setGenotype(treeToo.nodes, genotype.gene, genotype.positions, metadata.rootSequenceSecondTree);
+      }
     }
     const scaleType = genotype ? "categorical" : colorings[colorBy].type;
     if (genotype) {
-      ({legendValues, colorScale} = createScaleForGenotype(tree.nodes, genotype.aa));
+      ({legendValues, colorScale} = createScaleForGenotype(tree.nodes, treeToo?.nodes, genotype.aa));
       domain = [...legendValues];
     } else if (colorings && colorings[colorBy]) {
       if (scaleType === "continuous" || scaleType==="temporal") {
@@ -152,8 +155,8 @@ export function createNonContinuousScaleFromProvidedScaleMap(colorBy, providedSc
   };
 }
 
-function createScaleForGenotype(t1nodes, aaGenotype) {
-  const legendValues = orderOfGenotypeAppearance(t1nodes, aaGenotype);
+function createScaleForGenotype(t1nodes, t2nodes, aaGenotype) {
+  const legendValues = orderOfGenotypeAppearance(t1nodes, t2nodes, aaGenotype);
   const trueValues = aaGenotype ?
     legendValues.filter((x) => x !== "X" && x !== "-" && x !== "") :
     legendValues.filter((x) => x !== "X" && x !== "-" && x !== "N" && x !== "");

src/util/setGenotype.js

@@ -67,15 +67,20 @@ export const setGenotype = (nodes, prot, positions, rootSequence) => {
   // console.log(`set ${ancNodes.length} nodes to the ancestral state: ${ancState}`)
 };
 
-export const orderOfGenotypeAppearance = (nodes, aaGenotype) => {
+export const orderOfGenotypeAppearance = (nodes, nodesSecondTree, aaGenotype) => {
   const seen = {};
-  nodes.forEach((n) => {
+
+  function addGenotype(n) {
     let numDate = getTraitFromNode(n, "num_date");
     if (numDate === undefined) numDate = 0;
     if (!seen[n.currentGt] || numDate < seen[n.currentGt]) {
       seen[n.currentGt] = numDate;
     }
-  });
+  }
+
+  nodes.forEach(addGenotype);
+  if (nodesSecondTree) nodesSecondTree.forEach(addGenotype);
+
   const ordered = Object.keys(seen);
   ordered.sort((a, b) => seen[a] < seen[b] ? -1 : 1);
   let orderedBases;