From 495c5db3526734620811cc2d9c2aaea463fe383f Mon Sep 17 00:00:00 2001
From: Trevor Bedford <trevor@bedford.io>
Date: Wed, 15 Mar 2017 15:56:27 -0700
Subject: [PATCH] Map colors by genotype. Fixes #220.

---
 src/components/map/map.js     |  4 +++-
 src/util/mapHelpersLatLong.js | 12 +++++++-----
 src/util/treeHelpers.js       |  2 +-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/components/map/map.js b/src/components/map/map.js
index 3f26ba587..3b8427273 100644
--- a/src/components/map/map.js
+++ b/src/components/map/map.js
@@ -24,6 +24,7 @@ import getLatLongs from "../../util/mapHelpersLatLong";
     colorBy: state.controls.colorBy,
     map: state.map,
     geoResolution: state.controls.geoResolution,
+    sequences: state.sequences
   };
 })
 class Map extends React.Component {
@@ -225,7 +226,8 @@ class Map extends React.Component {
       this.state.map,
       this.props.colorBy,
       this.props.geoResolution,
-      this.props.colorScale.scale,
+      this.props.colorScale,
+      this.props.sequences
     );
   }
   createMap() {
diff --git a/src/util/mapHelpersLatLong.js b/src/util/mapHelpersLatLong.js
index 34ced162b..9b8c875ee 100644
--- a/src/util/mapHelpersLatLong.js
+++ b/src/util/mapHelpersLatLong.js
@@ -1,6 +1,7 @@
 import {averageColors} from "./colorHelpers";
+import {getTipColorAttribute} from "./treeHelpers";
 
-const getLatLongs = (nodes, metadata, map, colorBy, geoResolution, colorScale) => {
+const getLatLongs = (nodes, metadata, map, colorBy, geoResolution, colorScale, sequences) => {
 
   const aggregatedLocations = {}; /* demes */
   const aggregatedTransmissions = {}; /* edges, animation paths */
@@ -28,13 +29,14 @@ const getLatLongs = (nodes, metadata, map, colorBy, geoResolution, colorScale) =
     aggregate locations for demes
   */
   nodes.forEach((n) => {
+    const tipColorAttribute = getTipColorAttribute(n, colorScale, sequences);
     if (!n.children) {
       // look up geo1 geo2 geo3 do lat longs differ
       if (aggregatedLocations[n.attr[geoResolution]]) {
-        aggregatedLocations[n.attr[geoResolution]].push(colorScale(n.attr[colorBy]));
+        aggregatedLocations[n.attr[geoResolution]].push(colorScale.scale(tipColorAttribute));
       } else {
         // if we haven't added this pair, add it
-        aggregatedLocations[n.attr[geoResolution]] = [colorScale(n.attr[colorBy])];
+        aggregatedLocations[n.attr[geoResolution]] = [colorScale.scale(tipColorAttribute)];
       }
     }
     if (n.children) {
@@ -43,10 +45,10 @@ const getLatLongs = (nodes, metadata, map, colorBy, geoResolution, colorScale) =
         if (n.attr[geoResolution] === child.attr[geoResolution]) { return; }
         // look up in transmissions dictionary
         if (aggregatedTransmissions[n.attr[geoResolution] + "/" + child.attr[geoResolution]]) {
-          aggregatedTransmissions[n.attr[geoResolution] + "/" + child.attr[geoResolution]].push(colorScale(n.attr[colorBy]));
+          aggregatedTransmissions[n.attr[geoResolution] + "/" + child.attr[geoResolution]].push(colorScale.scale(tipColorAttribute));
         } else {
           // we don't have it, add it
-          aggregatedTransmissions[n.attr[geoResolution] + "/" + child.attr[geoResolution]] = [colorScale(n.attr[colorBy])];
+          aggregatedTransmissions[n.attr[geoResolution] + "/" + child.attr[geoResolution]] = [colorScale.scale(tipColorAttribute)];
         }
       });
     }
diff --git a/src/util/treeHelpers.js b/src/util/treeHelpers.js
index 5bfc34064..9665b4efd 100644
--- a/src/util/treeHelpers.js
+++ b/src/util/treeHelpers.js
@@ -136,7 +136,7 @@ export const calcBranchThickness = function (nodes, rootIdx) {
   return nodes.map((d) => freqScale(d.tipCount / maxTipCount));
 };
 
-const getTipColorAttribute = function (node, colorScale, sequences) {
+export const getTipColorAttribute = function (node, colorScale, sequences) {
   if (colorScale.colorBy.slice(0, 3) === "gt-" && colorScale.genotype) {
     return getGenotype(colorScale.genotype[0][0],
                        colorScale.genotype[0][1],