Add Color to Force Directed Graph

superset/assets/javascripts/explorev2/stores/fields.js

@@ -376,6 +376,23 @@ export const fields = {
                  'domain_granularity. Should be larger or equal to Time Grain',
   },
 
+  graph_color: {
+    type: 'SelectField',
+    label: 'Node Color',
+    default: [],
+    description: 'Choose a source node dimension to color the nodes.',
+    mapStateToProps: (state) => ({
+      choices: (state.datasource) ? state.datasource.all_cols : [],
+    }),
+  },
+
+  graph_labels: {
+    type: 'CheckboxField',
+    label: 'Node Lables',
+    default: true,
+    description: 'Display lables on all node?',
+  },
+
   link_length: {
     type: 'SelectField',
     freeForm: true,

superset/assets/javascripts/explorev2/stores/visTypes.js

@@ -544,6 +544,8 @@ const visTypes = {
       {
         label: 'Force Layout',
         fieldSetRows: [
+          ['graph_color'],
+          ['graph_labels'],
           ['link_length'],
           ['charge'],
         ],

superset/assets/visualizations/directed_force.js

@@ -10,6 +10,7 @@ const directedForceVis = function (slice, json) {
   const height = slice.height() - 25;
   const linkLength = json.form_data.link_length || 200;
   const charge = json.form_data.charge || -500;
+  const labels = json.form_data.graph_labels | 0;
 
   const links = json.data;
   const nodes = {};
@@ -46,6 +47,7 @@ const directedForceVis = function (slice, json) {
     }
 
     nodes[targetName].total += link.value;
+    nodes[targetName].color = link.color;
   });
 
   /* eslint-disable no-use-before-define */
@@ -124,11 +126,11 @@ const directedForceVis = function (slice, json) {
     .select('circle')
     .transition()
     .style('stroke-width', 5);
-
     d3.select(this)
     .select('text')
     .transition()
-    .style('font-size', 25);
+    .style('font-size', 25)
+    .style('opacity', 1);
   })
   .on('mouseleave', function () {
     d3.select(this)
@@ -138,7 +140,8 @@ const directedForceVis = function (slice, json) {
     d3.select(this)
     .select('text')
     .transition()
-    .style('font-size', 12);
+    .style('font-size', 12)
+    .style('opacity', labels);
   })
   .call(force.drag);
 
@@ -150,15 +153,21 @@ const directedForceVis = function (slice, json) {
   .domain(ext)
   .range([3, 30]);
 
+  const color = d3.scale.category20();
+
   node.append('circle')
   .attr('r', function (d) {
     return circleScale(Math.sqrt(d.total));
+  })
+  .style('fill', function (d) {
+    return color(d.color);
   });
 
   // add the text
   node.append('text')
   .attr('x', 6)
   .attr('dy', '.35em')
+  .style('opacity', labels)
   .text(function (d) {
     return d.name;
   });

superset/data/__init__.py

@@ -55,6 +55,7 @@ def load_energy():
         if_exists='replace',
         chunksize=500,
         dtype={
+            'cluster': String(255),
             'source': String(255),
             'target': String(255),
             'value': Float(),
@@ -123,6 +124,7 @@ def load_energy():
             ],
             "having": "",
             "link_length": "200",
+            "graph_color": "cluster",
             "metric": "sum__value",
             "row_limit": "5000",
             "slice_name": "Force",

superset/forms.py

@@ -321,6 +321,17 @@ def __init__(self, viz):
                 "choices": self.choicify(datasource.groupby_column_names),
                 "description": _("One or many fields to pivot as columns")
             }),
+            'graph_color': (SelectField, {
+                "label": _("Node Color"),
+                "default": None,
+                "choices": self.choicify([None] + datasource.column_names),
+                "description": _("Choose a source node dimension to color the nodes"),
+            }),
+            'graph_labels': (BetterBooleanField, {
+                "label": _("Node Labels"),
+                "default": True,
+                "description": _("Display labels on all node?")
+            }),
             'all_columns': (SelectMultipleSortableField, {
                 "label": _("Columns"),
                 "choices": self.choicify(datasource.column_names),

superset/viz.py

@@ -1759,6 +1759,8 @@ class DirectedForceViz(BaseViz):
     }, {
         'label': _('Force Layout'),
         'fields': (
+            'graph_color',
+            'graph_labels',
             'link_length',
             'charge',
         )
@@ -1775,11 +1777,16 @@ def query_obj(self):
         if len(self.form_data['groupby']) != 2:
             raise Exception("Pick exactly 2 columns to 'Group By'")
         qry['metrics'] = [self.form_data['metric']]
+        if self.form_data['graph_color'] != 'None':
+            qry['groupby'] += [self.form_data['graph_color']]
         return qry
 
     def get_data(self):
         df = self.get_df()
-        df.columns = ['source', 'target', 'value']
+        if self.form_data['graph_color'] == 'None':
+            df.columns = ['source', 'target', 'value']
+        else:
+            df.columns = ['source', 'target', 'color', 'value']
         return df.to_dict(orient='records')