Pass affinity to fix connectivity in linkage tree

brenolf · brenolf · commit 44554bdf9e86 · 2017-07-16T17:17:01.000-04:00
Resolves #9308.
diff --git a/sklearn/cluster/hierarchical.py b/sklearn/cluster/hierarchical.py
@@ -415,7 +415,9 @@ def linkage_tree(X, connectivity=None, n_components=None,
             return children_, 1, n_samples, None, distances
         return children_, 1, n_samples, None
 
-    connectivity, n_components = _fix_connectivity(X, connectivity)
+    connectivity, n_components = _fix_connectivity(X,
+                                                   connectivity,
+                                                   affinity=affinity)
 
     connectivity = connectivity.tocoo()
     # Put the diagonal to zero
diff --git a/sklearn/cluster/tests/test_hierarchical.py b/sklearn/cluster/tests/test_hierarchical.py
@@ -518,3 +518,32 @@ def test_agg_n_clusters():
         msg = ("n_clusters should be an integer greater than 0."
                " %s was provided." % str(agc.n_clusters))
         assert_raise_message(ValueError, msg, agc.fit, X)
+
+
+def test_affinity_passed_to_fix_connectivity():
+    # Test that the affinity parameter is actually passed to the pairwise
+    # function
+
+    rng = np.random.RandomState(0)
+    X = rng.randn(2, 2)
+
+    mask = np.zeros((2, 2), dtype=np.bool)
+    mask[0:1, 0:1] = True
+    mask[-1:, -1:] = True
+    mask = mask.reshape(2 ** 2)
+    connectivity = grid_to_graph(n_x=2, n_y=2, mask=mask, 
+                                 return_as=np.ndarray)
+
+    class FakeAffinity:
+      def __init__(self):
+        self.counter = 0
+
+      def increment(self, *args, **kwargs):
+        self.counter += 1
+        return self.counter
+
+    fa = FakeAffinity()
+
+    linkage_tree(X, affinity=fa.increment, connectivity=connectivity)
+
+    assert_equal(fa.counter, 3)
