Allow capping out of overplotting to speed up plots of condensed trees.

Author: lmcinnes

hdbscan/plots.py

@@ -53,7 +53,10 @@ def __init__(self, condensed_tree_array, cluster_selection_method='eom'):
         self._raw_tree = condensed_tree_array
         self.cluster_selection_method = cluster_selection_method
 
-    def get_plot_data(self, leaf_separation=1, log_size=False):
+    def get_plot_data(self,
+                      leaf_separation=1,
+                      log_size=False,
+                      max_rectangle_per_icicle=20):
         """Generates data for use in plotting the 'icicle plot' or dendrogram 
         plot of the condensed tree generated by HDBSCAN.
 
@@ -68,6 +71,12 @@ def get_plot_data(self, leaf_separation=1, log_size=False):
                    points in the cluster at a given lambda value).
                    (default False)
         
+        max_rectangles_per_icicle : int, optional
+            To simplify the plot this method will only emit
+            ``max_rectangles_per_icicle`` bars per branch of the dendrogram.
+            This ensures that we don't suffer from massive overplotting in
+            cases with a lot of data points.
+
         Returns
         -------
         plot_data : dict
@@ -127,22 +136,39 @@ def get_plot_data(self, leaf_separation=1, log_size=False):
             c_children = self._raw_tree[self._raw_tree['parent'] == c]
             current_size = np.sum(c_children['child_size'])
             current_lambda = cluster_y_coords[c]
+            cluster_max_size = current_size
+            cluster_max_lambda = c_children['lambda_val'].max()
+            cluster_min_size = np.sum(
+                c_children[c_children['lambda_val'] ==
+                           cluster_max_lambda]['child_size'])
 
             if log_size:
                 current_size = np.log(current_size)
+                cluster_max_size = np.log(cluster_max_size)
+                cluster_min_size = np.log(cluster_min_size)
+
+            total_size_change = float(cluster_max_size - cluster_min_size)
+            step_size_change = total_size_change / max_rectangle_per_icicle
 
             cluster_bounds[c][CB_LEFT] = cluster_x_coords[c] * scaling - (current_size / 2.0)
             cluster_bounds[c][CB_RIGHT] = cluster_x_coords[c] * scaling + (current_size / 2.0)
             cluster_bounds[c][CB_BOTTOM] = cluster_y_coords[c]
             cluster_bounds[c][CB_TOP] = np.max(c_children['lambda_val'])
 
+            last_step_size = current_size
+            last_step_lambda = current_lambda
+
             for i in np.argsort(c_children['lambda_val']):
                 row = c_children[i]
-                if row['lambda_val'] != current_lambda:
+                if row['lambda_val'] != current_lambda and \
+                        (last_step_size - current_size > step_size_change
+                        or row['lambda_val'] == cluster_max_lambda):
                     bar_centers.append(cluster_x_coords[c] * scaling)
-                    bar_tops.append(row['lambda_val'] - current_lambda)
-                    bar_bottoms.append(current_lambda)
-                    bar_widths.append(current_size)
+                    bar_tops.append(row['lambda_val'] - last_step_lambda)
+                    bar_bottoms.append(last_step_lambda)
+                    bar_widths.append(last_step_size)
+                    last_step_size = current_size
+                    last_step_lambda = current_lambda
                 if log_size:
                     exp_size = np.exp(current_size) - row['child_size']
                     # Ensure we don't try to take log of zero
@@ -184,6 +210,72 @@ def get_plot_data(self, leaf_separation=1, log_size=False):
             'cluster_bounds': cluster_bounds
         }
 
+    def get_simple_plot_data(self, leaf_separation=1, log_size=False,
+                             max_rectangle_per_icicle=20):
+        """Generates simplified data for use in plotting the 'icicle plot' or
+        dendrogram plot of the condensed tree generated by HDBSCAN.
+
+        Parameters
+        ----------
+        leaf_separation : float, optional
+                          How far apart to space the final leaves of the
+                          dendrogram. (default 1)
+
+        log_size : boolean, optional
+                   Use log scale for the 'size' of clusters (i.e. number of
+                   points in the cluster at a given lambda value).
+                   (default False)
+
+        max_rectangles_per_icicle : int, optional
+            To simplify the plot this method will only emit
+            ``max_rectangles_per_icicle`` bars per branch of the dendrogram.
+            This ensures that we don't suffer from massive overplotting in
+            cases with a lot of data points.
+
+        Returns
+        -------
+        plot_data : dict
+                    Data associated to bars in a bar plot:
+                        `bar_centers` x coordinate centers for bars
+                        `bar_tops` heights of bars in lambda scale
+                        `bar_bottoms` y coordinate of bottoms of bars
+                        `bar_widths` widths of the bars (in x coord scale)
+                        `bar_bounds` a 4-tuple of [left, right, bottom, top]
+                                     giving the bounds on a full set of
+                                     cluster bars
+                    Data associates with cluster splits:
+                        `line_xs` x coordinates for horiontal dendrogram lines
+                        `line_ys` y coordinates for horiontal dendrogram lines
+        """
+        leaves = _get_leaves(self._raw_tree)
+        last_leaf = self._raw_tree['parent'].max()
+        root = self._raw_tree['parent'].min()
+
+        # We want to get the x and y coordinates for the start of each cluster
+        # Initialize the leaves, since we know where they go, the iterate
+        # through everything from the leaves back, setting coords as we go
+        cluster_x_coords = dict(zip(leaves, [leaf_separation * x
+                                             for x in range(len(leaves))]))
+        cluster_y_coords = {root: 0.0}
+
+        # We want to get the x and y coordinates for the start of each cluster
+        # Initialize the leaves, since we know where they go, the iterate
+        # through everything from the leaves back, setting coords as we go
+        cluster_x_coords = dict(zip(leaves, [leaf_separation * x
+                                             for x in range(len(leaves))]))
+        cluster_y_coords = {root: 0.0}
+
+        for cluster in range(last_leaf, root - 1, -1):
+            split = self._raw_tree[['child', 'lambda_val']]
+            split = split[(self._raw_tree['parent'] == cluster) &
+                          (self._raw_tree['child_size'] > 1)]
+            if len(split['child']) > 1:
+                left_child, right_child = split['child']
+                cluster_x_coords[cluster] = np.mean([cluster_x_coords[left_child],
+                                                     cluster_x_coords[right_child]])
+                cluster_y_coords[left_child] = split['lambda_val'][0]
+                cluster_y_coords[right_child] = split['lambda_val'][1]
+
     def _select_clusters(self):
         if self.cluster_selection_method == 'eom':
             stability = compute_stability(self._raw_tree)
@@ -213,7 +305,8 @@ def _select_clusters(self):
 
     def plot(self, leaf_separation=1, cmap='viridis', select_clusters=False,
              label_clusters=False, selection_palette=None,
-             axis=None, colorbar=True, log_size=False):
+             axis=None, colorbar=True, log_size=False,
+             max_rectangles_per_icicle=20):
         """Use matplotlib to plot an 'icicle plot' dendrogram of the condensed tree.
 
         Effectively this is a dendrogram where the width of each cluster bar is
@@ -224,45 +317,51 @@ def plot(self, leaf_separation=1, cmap='viridis', select_clusters=False,
 
         Parameters
         ----------
-        leaf_separation : float, optional
+        leaf_separation : float, optional (default 1)
                           How far apart to space the final leaves of the 
-                          dendrogram. (default 1)
+                          dendrogram.
 
-        cmap : string or matplotlib colormap, optional
+        cmap : string or matplotlib colormap, optional (default viridis)
                The matplotlib colormap to use to color the cluster bars.
-               (default viridis)
 
-        select_clusters : boolean, optional
+
+        select_clusters : boolean, optional (default False)
                           Whether to draw ovals highlighting which cluster
                           bar represent the clusters that were selected by
-                          HDBSCAN as the final clusters. (default False)
+                          HDBSCAN as the final clusters.
 
-        label_clusters : boolean, optional
+        label_clusters : boolean, optional (default False)
                          If select_clusters is True then this determines
                          whether to draw text labels on the clusters.
 
-        selection_palette : list of colors, optional
+        selection_palette : list of colors, optional (default None)
                             If not None, and at least as long as
                             the number of clusters, draw ovals
                             in colors iterating through this palette.
                             This can aid in cluster identification
                             when plotting.
 
-        axis : matplotlib axis or None, optional
+        axis : matplotlib axis or None, optional (default None)
                The matplotlib axis to render to. If None then a new axis
                will be generated. The rendered axis will be returned.
-               (default None)
 
-        colorbar : boolean, optional
+
+        colorbar : boolean, optional (default True)
                    Whether to draw a matplotlib colorbar displaying the range
-                   of cluster sizes as per the colormap. (default True)
+                   of cluster sizes as per the colormap.
 
-        log_size : boolean, optional
+        log_size : boolean, optional (default False)
                    Use log scale for the 'size' of clusters (i.e. number of
                    points in the cluster at a given lambda value).
-                   (default False)
+
            
-        Returns
+        max_rectangles_per_icicle : int, optional (default 20)
+            To simplify the plot this method will only emit
+            ``max_rectangles_per_icicle`` bars per branch of the dendrogram.
+            This ensures that we don't suffer from massive overplotting in
+            cases with a lot of data points.
+
+         Returns
         -------
         axis : matplotlib axis
                The axis on which the 'icicle plot' has been rendered.
@@ -274,7 +373,9 @@ def plot(self, leaf_separation=1, cmap='viridis', select_clusters=False,
                 'You must install the matplotlib library to plot the condensed tree.'
                 'Use get_plot_data to calculate the relevant data without plotting.')
 
-        plot_data = self.get_plot_data(leaf_separation=leaf_separation, log_size=log_size)
+        plot_data = self.get_plot_data(leaf_separation=leaf_separation,
+                                       log_size=log_size,
+                                       max_rectangle_per_icicle=max_rectangles_per_icicle)
 
         if cmap != 'none':
             sm = plt.cm.ScalarMappable(cmap=cmap,