From a86a5398ec11a4f66567e3073e2864ec8b87e725 Mon Sep 17 00:00:00 2001
From: Colin Megill <colinmegill@gmail.com>
Date: Fri, 16 Apr 2021 01:59:20 -0700
Subject: [PATCH] color graph by mean expression

---
 client/src/components/graph/graph.js         |  7 ++++--
 client/src/util/stateManager/colorHelpers.js | 26 +++++++++++++++++++-
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/client/src/components/graph/graph.js b/client/src/components/graph/graph.js
index 86ccc2b3..75571d8b 100644
--- a/client/src/components/graph/graph.js
+++ b/client/src/components/graph/graph.js
@@ -76,6 +76,7 @@ function createModelTF() {
   graphInteractionMode: state.controls.graphInteractionMode,
   colors: state.colors,
   pointDilation: state.pointDilation,
+  genesets: state.genesets.genesets,
 }))
 class Graph extends React.Component {
   static createReglState(canvas) {
@@ -526,6 +527,7 @@ class Graph extends React.Component {
       colorsProp,
       pointDilation
     );
+
     const { currentDimNames } = layoutChoice;
     const X = layoutDf.col(currentDimNames[0]).asArray();
     const Y = layoutDf.col(currentDimNames[1]).asArray();
@@ -782,10 +784,11 @@ class Graph extends React.Component {
   }
 
   createColorByQuery(colors) {
-    const { annoMatrix } = this.props;
+    const { annoMatrix, genesets } = this.props;
     const { schema } = annoMatrix;
     const { colorMode, colorAccessor } = colors;
-    return createColorQuery(colorMode, colorAccessor, schema);
+
+    return createColorQuery(colorMode, colorAccessor, schema, genesets);
   }
 
   renderPoints(
diff --git a/client/src/util/stateManager/colorHelpers.js b/client/src/util/stateManager/colorHelpers.js
index 4692d48c..f8fc360f 100644
--- a/client/src/util/stateManager/colorHelpers.js
+++ b/client/src/util/stateManager/colorHelpers.js
@@ -12,7 +12,7 @@ import { range } from "../range";
 given a color mode & accessor, generate an annoMatrix query that will
 fulfill it
 */
-export function createColorQuery(colorMode, colorByAccessor, schema) {
+export function createColorQuery(colorMode, colorByAccessor, schema, genesets) {
   if (!colorMode || !colorByAccessor || !schema) return null;
   switch (colorMode) {
     case "color by categorical metadata":
@@ -33,6 +33,25 @@ export function createColorQuery(colorMode, colorByAccessor, schema) {
         },
       ];
     }
+    case "color by geneset mean expression": {
+      const varIndex = schema?.annotations?.var?.index;
+      if (!varIndex) return null;
+
+      const _geneset = genesets.get(colorByAccessor);
+      const _setGenes = Array.from(_geneset.genes.keys());
+
+      return [
+        "X",
+        {
+          summarize: {
+            method: "mean",
+            field: "var",
+            column: varIndex,
+            values: _setGenes,
+          },
+        },
+      ];
+    }
     default: {
       return null;
     }
@@ -86,6 +105,11 @@ function _createColorTable(
       const { min, max } = col.summarize();
       return createColorsByContinuousMetadata(col.asArray(), min, max);
     }
+    case "color by geneset mean expression": {
+      const col = colorByData.icol(0);
+      const { min, max } = col.summarize();
+      return createColorsByContinuousMetadata(col.asArray(), min, max);
+    }
     default: {
       return defaultColors(schema.dataframe.nObs);
     }