Allow scale information to be supplied in an auspice config file

Auspice 2.26.0 can now use custom colors for continuous scales via the `scale` property on such a coloring; previously custom colors were only available for discrete scales (see  nextstrain/auspice#1340 for more).

Here we extend the schema for the auspice config JSON and update `augur export` to allow export of such scale information for continuous colorings.

Previously, scale information (for discrete scales) could only come from a colors TSV file supplied via `augur export --colors <TSV>`. We now allow this to be specified in the auspice config JSON, and use this information preferentially over the TSV.

augur/data/schema-auspice-config-v2.json

@@ -35,6 +35,21 @@
                         "type": "string",
                         "enum": ["continuous", "ordinal", "categorical", "boolean"]
                     },
+                    "scale": {
+                        "description": "Provided mapping between trait values & hex values",
+                        "$comment": "NOTE: if supplied here, we will not use information supplied to `augur export` via `--colors` for this coloring.",
+                        "type": "array",
+                        "items": {
+                            "type": "array",
+                            "items": [
+                                {
+                                    "type": ["string", "number"],
+                                    "description": "For categorical/ordinal scales, this is the (string) value of the trait to associate with the colour. For continuous scales this is the (numeric) value to associate to with the colour, and interpolation will be used to span the domain"
+                                },
+                                {"type": "string", "description": "color hex value", "pattern": "^#[0-9A-Fa-f]{6}$"}
+                            ]
+                        }
+                    },
                     "legend": {
                         "description": "Specify the entries displayed in the legend. This can be used to restrict the entries in the legend for display without otherwise affecting the data viz",
                         "type": "array",

augur/data/schema-export-v2.json

@@ -180,11 +180,15 @@
                             },
                             "scale": {
                                 "description": "Provided mapping between trait values & hex values",
+                                "$comment": "For continuous scales at least 2 items must be specified",
                                 "type": "array",
                                 "items": {
                                     "type": "array",
                                     "items": [
-                                        {"type": "string", "description": "value of trait (should exist on >= 1 nodes)"},
+                                        {
+                                            "type": ["string", "number"],
+                                            "description": "For categorical/ordinal scales, this is the (string) value of the trait to associate with the colour. For continuous scales this is the (numeric) value to associate to with the colour, and interpolation will be used to span the domain"
+                                        },
                                         {"type": "string", "description": "color hex value", "pattern": "^#[0-9A-Fa-f]{6}$"}
                                     ]
                                 }

augur/export_v2.py

@@ -6,6 +6,7 @@
 import time
 from collections import defaultdict, deque
 import warnings
+import numbers
 import re
 from Bio import Phylo
 from .utils import read_metadata, read_node_data, write_json, read_config, read_lat_longs, read_colors
@@ -211,8 +212,32 @@ def _get_title(key):
         return key
 
     def _add_color_scale(coloring):
+        ## consider various sources to find any user-provided scale information
         key = coloring["key"]
-        if key.lower() in provided_colors:
+        scale_type = coloring["type"]
+        if scale_type is "continuous":
+            ## continuous scale information can only come from an auspice config JSON
+            if config.get(key, {}).get("scale"):
+                # enforce numeric values (we can't use the schema for this)
+                provided_scale = [s for s in config[key]['scale'] if isinstance(s[0], numbers.Number)]
+                if len(provided_scale)<2:
+                    warn(f"The scale provided for {key} had fewer than 2 (valid numeric) entries. Skipping.")
+                    return coloring
+                coloring["scale"] = provided_scale
+                return coloring
+        elif config.get(key, {}).get("scale"):
+            # we preferentially use a scale explicitly provided in the auspice config JSON
+            values_in_tree = get_values_across_nodes(node_attrs, key)
+            scale = [info for info in config[key]['scale'] if info[0] in values_in_tree]
+            provided_values_unseen_in_tree = [info[0] for info in config[key]['scale'] if info[0] not in values_in_tree]
+            if len(scale):
+                coloring["scale"] = scale
+                if len(provided_values_unseen_in_tree):
+                    warn(f"These values for \"{key}\" were not specified in the scale your configuration file provided:\n\t{', '.join(provided_values_unseen_in_tree)}.\n\tAuspice will create colors for them.")
+                return coloring
+            warn(f"The configuration JSON specifies a color scale for {key} however none of the values in the tree are in this scale! Auspice will generate its own color scale for this trait.")
+        elif key.lower() in provided_colors:
+            # `provided_colors` typically originates from a colors.tsv file
             scale = []
             trait_values = {str(val).lower(): val for val in get_values_across_nodes(node_attrs, key)}
             trait_values_unseen = {k for k in trait_values}
@@ -223,9 +248,10 @@ def _add_color_scale(coloring):
             if len(scale):
                 coloring["scale"] = scale
                 if len(trait_values_unseen):
-                    warn("These values for trait {} were not specified in your provided color scale: {}. Auspice will create colors for them.".format(key, ", ".join(trait_values_unseen)))
-            else:
-                warn("You've specified a color scale for {} but none of the values found on the tree had associated colors. Auspice will generate its own color scale for this trait.".format(key))
+                    warn(f"These values for trait {key} were not specified in the colors file you provided:\n\t{', '.join(trait_values_unseen)}.\n\tAuspice will create colors for them.")
+                return coloring
+            warn(f"You've supplied a colors file with information for {key} but none of the values found on the tree had associated colors. Auspice will generate its own color scale for this trait.")
+        # Fallthrough (no scale information provided means that auspice will generate its own scale)
         return coloring
 
     def _add_legend(coloring):