Improve plotting coverage

lmcinnes · Aug 20, 2016 · 32ba543 · 32ba543
1 parent 4243ef1
commit 32ba543
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/hdbscan/plots.py b/hdbscan/plots.py
@@ -649,7 +649,7 @@ def __init__(self, mst, data):
 
     def plot(self, axis=None, node_size=40, node_color='k',
              node_alpha=0.8, edge_alpha=0.5, edge_cmap='viridis_r',
-             edge_linewidth=2, vary_linewidth=True, colorbar=True):
+             edge_linewidth=2, vary_line_width=True, colorbar=True):
         """Plot the minimum spanning tree (as projected into 2D by t-SNE if required).
 
         Parameters
@@ -680,7 +680,7 @@ def plot(self, axis=None, node_size=40, node_color='k',
         edge_linewidth : float, optional
                 The linewidth to use for rendering edges (default 2).
 
-        vary_linewidth : bool, optional
+        vary_line_width : bool, optional
                 Edge width is proportional to (log of) the inverse of the
                 mutual reachability distance. (default True)
 
@@ -718,7 +718,7 @@ def plot(self, axis=None, node_size=40, node_color='k',
         else:
             projection = self._data.copy()
 
-        if vary_linewidth:
+        if vary_line_width:
             line_width = edge_linewidth * (np.log(self._mst.T[2].max() / self._mst.T[2]) + 1.0)
         else:
             line_width = edge_linewidth

diff --git a/hdbscan/tests/test_hdbscan.py b/hdbscan/tests/test_hdbscan.py
@@ -271,9 +271,37 @@ def test_hdbscan_boruvka_balltree_matches():
     assert_less(num_mismatches / float(data.shape[0]), 0.15)
 
 def test_condensed_tree_plot():
-    clusterer = HDBSCAN().fit(X)
-    clusterer.condensed_tree_.get_plot_data()
-    if_matplotlib(clusterer.condensed_tree_.plot)(select_clusters=True, selection_palette=('r','g','b'))
+    clusterer = HDBSCAN(gen_min_span_tree=True).fit(X)
+    if_matplotlib(clusterer.condensed_tree_.plot)(select_clusters=True,
+                                                  selection_palette=('r','g','b'),
+                                                  cmap='Reds')
+    if_matplotlib(clusterer.condensed_tree_.plot)(label_clusters=True,
+                                                  colorbar=False,
+                                                  cmap='none')
+
+def test_single_linkage_tree_plot():
+    clusterer = HDBSCAN(gen_min_span_tree=True).fit(X)
+    if_matplotlib(clusterer.single_linkage_tree_.plot)()
+    if_matplotlib(clusterer.single_linkage_tree_.plot)(vary_line_width=False,
+                                                       truncate_mode='lastp',
+                                                       p=10,
+                                                       colorbar=False)
+
+def test_min_span_tree_plot():
+    clusterer = HDBSCAN(gen_min_span_tree=True).fit(X)
+    if_matplotlib(clusterer.minimum_spanning_tree_.plot)()
+
+    H, y = make_blobs(n_samples=50, random_state=0, n_features=10)
+    H = StandardScaler().fit_transform(H)
+
+    clusterer = HDBSCAN(gen_min_span_tree=True).fit(H)
+    if_matplotlib(clusterer.minimum_spanning_tree_.plot)(vary_line_width=False, colorbar=False)
+
+    H, y = make_blobs(n_samples=50, random_state=0, n_features=40)
+    H = StandardScaler().fit_transform(H)
+
+    clusterer = HDBSCAN(gen_min_span_tree=True).fit(H)
+    if_matplotlib(clusterer.minimum_spanning_tree_.plot)(vary_line_width=False, colorbar=False)
 
 def test_tree_numpy_output_formats():