From 95a83ddeac8218b97f6d22a186b8e72bdee232ce Mon Sep 17 00:00:00 2001
From: Eric Ma <ericmajinglong@gmail.com>
Date: Mon, 18 Sep 2017 09:56:37 -0400
Subject: [PATCH] Added auto-colorbar. New version good to go.

---
 HISTORY.rst     |  6 ++++++
 README.rst      | 10 +++++++++-
 environment.yml |  1 +
 nxviz/plots.py  | 41 ++++++++++++++++++++++++++++++++++-------
 setup.py        |  2 +-
 5 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index 5abcac4d..e0d94de9 100755
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -2,6 +2,12 @@
 History
 =======
 
+0.3.2 (2017-09-18)
+------------------
+
+* All plots except for HivePlot are implemented.
+* Implemented auto-colorbar for plots that have continuous node colors.
+
 0.1.0 (2016-07-15)
 ------------------
 
diff --git a/README.rst b/README.rst
index 924e64a6..4e644cda 100644
--- a/README.rst
+++ b/README.rst
@@ -37,13 +37,21 @@ nxviz
 
 * This is free software distributed under the MIT License.
 
-
 Features
 --------
 
 * Declarative API.
 * Works with NetworkX, one of the more popular graph libraries in Python.
 
+Feature Requests
+----------------
+
+If you have a feature request, please post it as an issue on the GitHub repository issue_ tracker. Even better, put in a PR for it! I am more than happy to guide you through the codebase so that you can put in a contribution to the codebase - and I'll give you a digital `nxviz` contributor badge that you can put on your personal website, as a way of saying thanks!
+
+Because nxviz is currently maintained by volunteers and has no fiscal support, any feature requests will be prioritized according to what maintainers encounter as a need in our day-to-day jobs. Please temper expectations accordingly.
+
+.. _issue: https://github.com/ericmjl/nxviz/issues
+
 Credits
 ---------
 
diff --git a/environment.yml b/environment.yml
index 4a8daac2..f849716e 100644
--- a/environment.yml
+++ b/environment.yml
@@ -11,6 +11,7 @@ dependencies:
 - pip
 - setuptools
 - pycodestyle
+- pydocstyle
 - click
 - pandas
 - coverage
diff --git a/nxviz/plots.py b/nxviz/plots.py
index 5dc36f22..45cba7e0 100644
--- a/nxviz/plots.py
+++ b/nxviz/plots.py
@@ -9,6 +9,10 @@
 from .utils import (cmaps, infer_data_type, is_data_diverging,
                     num_discrete_groups)
 
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
 
 def despine(ax):
     for spine in ax.spines:
@@ -79,6 +83,8 @@ def __init__(self, graph, node_order=None, node_size=None,
 
         # Set node colors
         self.node_color = node_color
+        self.sm = None  # sm -> for scalarmappable. See https://stackoverflow.com/questions/8342549/matplotlib-add-colorbar-to-a-sequence-of-line-plots  # noqa
+        logging.debug(f'INIT: {self.sm}')
         if self.node_color:
             self.node_colors = []
             self.compute_node_colors()
@@ -132,16 +138,25 @@ def check_data_types(self, data_types):
     def draw(self):
         """
         Draws the Plot to screen.
+
+        If there is a continuous datatype for the nodes, it will be reflected
+        in self.sm being constructed (in `compute_node_colors`). It will then
+        automatically add in a colorbar to the plot and scale the plot axes
+        accordingly.
         """
         self.draw_nodes()
         self.draw_edges()
+        logging.debug(f'DRAW: {self.sm}')
+        if self.sm:
+            self.figure.subplots_adjust(right=0.8)
+            cax = self.figure.add_axes([0.85, 0.2, 0.05, 0.6])
+            self.figure.colorbar(self.sm, cax=cax)
         self.ax.relim()
         self.ax.autoscale_view()
+        self.ax.set_aspect('equal')
 
     def compute_node_colors(self):
-        """
-        Computes the node colors.
-        """
+        """Compute the node colors. Also computes the colorbar."""
         data = [self.graph.node[n][self.node_color] for n in self.nodes]
         data_reduced = sorted(list(set(data)))
         dtype = infer_data_type(data)
@@ -158,6 +173,17 @@ def compute_node_colors(self):
             idx = data_reduced.index(d) / n_grps
             self.node_colors.append(cmap(idx))
 
+        # Add colorbar if required.
+        logging.debug(f'length of data_reduced: {len(data_reduced)}')
+        logging.debug(f'dtype: {dtype}')
+        if len(data_reduced) > 1 and dtype == 'continuous':
+            self.sm = plt.cm.ScalarMappable(cmap=cmap,
+                                            norm=plt.Normalize(vmin=min(data_reduced),  # noqa
+                                                               vmax=max(data_reduced)   # noqa
+                                                               )
+                                            )
+            self.sm._A = []
+
     def compute_node_positions(self):
         """
         Computes the positions of each node on the plot.
@@ -436,6 +462,10 @@ def __init__(self, graph, node_order=None, node_size=None,
     def draw(self):
         """
         Draws the plot to screen.
+
+        Note to self: Do NOT call super(MatrixPlot, self).draw(); the
+        underlying logic for drawing here is completely different from other
+        plots, and as such necessitates a different implementation.
         """
         matrix = nx.to_numpy_matrix(self.graph, nodelist=self.nodes)
         self.ax.matshow(matrix, cmap=self.cmap)
@@ -517,10 +547,7 @@ def draw_edges(self):
             self.ax.add_patch(patch)
 
     def draw(self):
-        self.draw_nodes()
-        self.draw_edges()
+        super(ArcPlot, self).draw()
         xlimits = (-1, len(self.nodes) + 1)
-        # halfwidth = len(self.nodes) + 1 / 2
-        # ylimits = (-halfwidth, halfwidth)
         self.ax.set_xlim(*xlimits)
         self.ax.set_ylim(*xlimits)
diff --git a/setup.py b/setup.py
index b927b506..b7755425 100644
--- a/setup.py
+++ b/setup.py
@@ -13,7 +13,7 @@
 
 setup(
     name='nxviz',
-    version="0.3.1",
+    version="0.3.2",
     description="Graph Visualization Package",
     long_description=readme + '\n\n' + history,
     author="Eric J. Ma",