Added color_threshold to dendrogram

Author: paulamool

plotly/figure_factory/_dendrogram.py

@@ -17,7 +17,7 @@
 def create_dendrogram(X, orientation="bottom", labels=None,
                       colorscale=None, distfun=None,
                       linkagefun=lambda x: sch.linkage(x, 'complete'),
-                      hovertext=None):
+                      hovertext=None, color_threshold=None):
     """
     BETA function that returns a dendrogram Plotly figure object.
 
@@ -32,6 +32,7 @@ def create_dendrogram(X, orientation="bottom", labels=None,
     :param (list[list]) hovertext: List of hovertext for constituent traces of dendrogram
 
         clusters
+    :param (function) color_threshold: function to colour clusters under a set threshold
 
     Example 1: Simple bottom oriented dendrogram
     ```
@@ -88,7 +89,7 @@ def create_dendrogram(X, orientation="bottom", labels=None,
 
     dendrogram = _Dendrogram(X, orientation, labels, colorscale,
                              distfun=distfun, linkagefun=linkagefun,
-                             hovertext=hovertext)
+                             hovertext=hovertext, color_threshold=color_threshold)
 
     return graph_objs.Figure(data=dendrogram.data, layout=dendrogram.layout)
 
@@ -100,7 +101,7 @@ def __init__(self, X, orientation='bottom', labels=None, colorscale=None,
                  width="100%", height="100%", xaxis='xaxis', yaxis='yaxis',
                  distfun=None,
                  linkagefun=lambda x: sch.linkage(x, 'complete'),
-                 hovertext=None):
+                 hovertext=None, color_threshold=None):
         self.orientation = orientation
         self.labels = labels
         self.xaxis = xaxis
@@ -127,7 +128,8 @@ def __init__(self, X, orientation='bottom', labels=None, colorscale=None,
             ordered_labels, leaves) = self.get_dendrogram_traces(X, colorscale,
                                                                  distfun,
                                                                  linkagefun, 
-                                                                 hovertext)
+                                                                 hovertext,
+                                                                 color_threshold)
 
         self.labels = ordered_labels
         self.leaves = leaves
@@ -236,7 +238,7 @@ def set_figure_layout(self, width, height):
 
         return self.layout
 
-    def get_dendrogram_traces(self, X, colorscale, distfun, linkagefun, hovertext):
+    def get_dendrogram_traces(self, X, colorscale, distfun, linkagefun, hovertext, color_threshold):
         """
         Calculates all the elements needed for plotting a dendrogram.
 
@@ -247,6 +249,7 @@ def get_dendrogram_traces(self, X, colorscale, distfun, linkagefun, hovertext):
         :param (function) linkagefun: Function to compute the linkage matrix
                                       from the pairwise distances
         :param (list) hovertext: List of hovertext for constituent traces of dendrogram
+        :param (list) color_threshold: sets the color of cluster under set threshold
         :rtype (tuple): Contains all the traces in the following order:
             (a) trace_list: List of Plotly trace objects for dendrogram tree
             (b) icoord: All X points of the dendrogram tree as array of arrays
@@ -261,7 +264,7 @@ def get_dendrogram_traces(self, X, colorscale, distfun, linkagefun, hovertext):
         d = distfun(X)
         Z = linkagefun(d)
         P = sch.dendrogram(Z, orientation=self.orientation,
-                           labels=self.labels, no_plot=True)
+                           labels=self.labels, no_plot=True, color_threshold=color_threshold)
 
         icoord = scp.array(P['icoord'])
         dcoord = scp.array(P['dcoord'])