From b29f6e48d0c077a6c216674fd2a1a8563661ba25 Mon Sep 17 00:00:00 2001
From: james hadfield <hadfield.james@gmail.com>
Date: Tue, 27 Aug 2024 15:24:08 +1200
Subject: [PATCH] Allow YYYY-MM-DD values in temporal color scale

These values may be ambiguous (using "XX" notation). There is no ability
to provide a custom scale / legend for temporal scales, with errors
printed to the console if these are attempted.
---
 src/util/colorHelpers.js     | 17 +++++++++++++++++
 src/util/colorScale.js       | 24 ++++++++++++++++++------
 src/util/tipRadiusHelpers.ts |  7 +++++--
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/src/util/colorHelpers.js b/src/util/colorHelpers.js
index d202c2bdd..8885e3a40 100644
--- a/src/util/colorHelpers.js
+++ b/src/util/colorHelpers.js
@@ -4,6 +4,7 @@ import scalePow from "d3-scale/src/pow";
 import { isColorByGenotype, decodeColorByGenotype } from "./getGenotype";
 import { getTraitFromNode } from "./treeMiscHelpers";
 import { isValueValid } from "./globals";
+import { calendarToNumeric } from "./dateHelpers";
 
 /**
  * Average over the visible colours for a given location
@@ -147,3 +148,19 @@ export const getColorByTitle = (colorings, colorBy) => {
   return colorings[colorBy] === undefined ?
     "" : colorings[colorBy].title;
 };
+
+/**
+ * We allow values (on nodes) to be encoded as numeric dates (2021.123) or
+ * YYYY-MM-DD strings. This helper function handles this flexibility and
+ * translates any provided value to either a number or undefined.
+ */
+export function numDate(value) {
+  switch (typeof value) {
+    case "number":
+      return value;
+    case "string":
+      return calendarToNumeric(value, true); // allow XX ambiguity
+    default:
+      return undefined;
+  }
+}
\ No newline at end of file
diff --git a/src/util/colorScale.js b/src/util/colorScale.js
index 0e745363b..76c4372a4 100644
--- a/src/util/colorScale.js
+++ b/src/util/colorScale.js
@@ -5,7 +5,7 @@ import { rgb } from "d3-color";
 import { interpolateHcl } from "d3-interpolate";
 import { genericDomain, colors, genotypeColors, isValueValid, NODE_VISIBLE } from "./globals";
 import { countTraitsAcrossTree } from "./treeCountingHelpers";
-import { getExtraVals } from "./colorHelpers";
+import { getExtraVals, numDate } from "./colorHelpers";
 import { isColorByGenotype, decodeColorByGenotype } from "./getGenotype";
 import { setGenotype, orderOfGenotypeAppearance } from "./setGenotype";
 import { getTraitFromNode } from "./treeMiscHelpers";
@@ -274,19 +274,22 @@ function createTemporalScale(colorBy, providedScale, t1nodes, t2nodes) {
   /* user-defined anchor points across the scale - note previously temporal scales could use this,
   although I doubt any did */
   // const anchorPoints = _validateContinuousAnchorPoints(providedScale);
+  if (providedScale) {
+    console.error("Auspice currently doesn't allow a JSON-provided 'scale' for temporal colorings");
+  }
 
   /* construct a domain / range which "focuses" on the tip dates, and be spaced according to sampling */
   let domain, range;
   {
-    let rootDate = getTraitFromNode(t1nodes[0], colorBy);
+    let rootDate = numDate(getTraitFromNode(t1nodes[0], colorBy));
     let vals = t1nodes.filter((n) => !n.hasChildren)
-      .map((n) => getTraitFromNode(n, colorBy));
+      .map((n) => numDate(getTraitFromNode(n, colorBy)));
     if (t2nodes) {
-      const treeTooRootDate = getTraitFromNode(t2nodes[0], colorBy);
+      const treeTooRootDate = numDate(getTraitFromNode(t2nodes[0], colorBy));
       if (treeTooRootDate < rootDate) rootDate = treeTooRootDate;
       vals.concat(
         t2nodes.filter((n) => !n.hasChildren)
-          .map((n) => getTraitFromNode(n, colorBy))
+          .map((n) => numDate(getTraitFromNode(n, colorBy)))
       );
     }
     vals = vals.sort();
@@ -306,9 +309,14 @@ function createTemporalScale(colorBy, providedScale, t1nodes, t2nodes) {
   // Hack to avoid a bug: https://github.com/nextstrain/auspice/issues/540
   if (Object.is(legendValues[0], -0)) legendValues[0] = 0;
 
+  const colorScale = (val) => {
+    const d = numDate(val);
+    return d===undefined ? unknownColor : scale(d);
+  };
+
   return {
     continuous: true,
-    colorScale: (val) => isValueValid(val) ? scale(val) : unknownColor,
+    colorScale,
     legendBounds: createLegendBounds(legendValues),
     legendValues
   };
@@ -457,6 +465,10 @@ function _validateContinuousAnchorPoints(providedScale) {
  */
 function parseUserProvidedLegendData(providedLegend, currentLegendValues, scaleType) {
   if (!Array.isArray(providedLegend)) return false;
+  if (scaleType==='temporal') {
+    console.error("Auspice currently doesn't allow a JSON-provided 'legend' for temporal colorings");
+    return false;
+  }
 
   const data = scaleType==="continuous" ?
     providedLegend.filter((d) => typeof d.value === "number") : // continuous scales _must_ have numeric stops
diff --git a/src/util/tipRadiusHelpers.ts b/src/util/tipRadiusHelpers.ts
index 7838af958..df612fc27 100644
--- a/src/util/tipRadiusHelpers.ts
+++ b/src/util/tipRadiusHelpers.ts
@@ -1,5 +1,5 @@
 import { tipRadius, tipRadiusOnLegendMatch } from "./globals";
-import { getTipColorAttribute } from "./colorHelpers";
+import { getTipColorAttribute, numDate } from "./colorHelpers";
 import { getTraitFromNode } from "./treeMiscHelpers";
 
 /**
@@ -13,7 +13,10 @@ import { getTraitFromNode } from "./treeMiscHelpers";
 * @returns bool
 */
 export const determineLegendMatch = (selectedLegendItem: (string|number), node:any, colorScale:any) => {
-  const nodeAttr = getTipColorAttribute(node, colorScale);
+  let nodeAttr = getTipColorAttribute(node, colorScale);
+  if (colorScale.scaleType === 'temporal') {
+    nodeAttr = numDate(nodeAttr);
+  }
   if (colorScale.continuous) {
     if (selectedLegendItem === colorScale.legendValues[0] && nodeAttr===colorScale.legendBounds[selectedLegendItem][0]) {
       return true;