callable bandwidth should operate on the rows of the distance matrix

Author: scottgigante

graphtools/api.py

@@ -65,10 +65,11 @@ def Graph(data,
     decay : `int` or `None`, optional (default: 10)
         Rate of alpha decay to use. If `None`, alpha decay is not used.
 
-    bandwidth : `float`, list-like or `None`, optional (default: `None`)
+    bandwidth : `float`, list-like,`callable`, or `None`, optional (default: `None`)
         Fixed bandwidth to use. If given, overrides `knn`. Can be a single
-        bandwidth or a list-like (shape=[n_samples]) of bandwidths for each
-        sample.
+        bandwidth, list-like (shape=[n_samples]) of bandwidths for each
+        sample, or a `callable` that takes in a `n x m` matrix and returns a
+        a single value or list-like of length n (shape=[n_samples])
 
     bandwidth_scale : `float`, optional (default : 1.0)
         Rescaling factor for bandwidth.

graphtools/graphs.py

@@ -721,8 +721,8 @@ class TraditionalGraph(DataGraph):
     bandwidth : `float`, list-like,`callable`, or `None`, optional (default: `None`)
         Fixed bandwidth to use. If given, overrides `knn`. Can be a single
         bandwidth, list-like (shape=[n_samples]) of bandwidths for each
-        sample, or a `callable` that takes in a square matrix and returns a
-        a single value or list-like(shape=[n_samples])
+        sample, or a `callable` that takes in a `n x m` matrix and returns a
+        a single value or list-like of length n (shape=[n_samples])
 
     bandwidth_scale : `float`, optional (default : 1.0)
         Rescaling factor for bandwidth.

test/test_exact.py

@@ -469,7 +469,7 @@ def test_build_dense_exact_kernel_to_data(**kwargs):
 
 
 def test_build_dense_exact_callable_bw_kernel_to_data(**kwargs):
-    G = build_graph(data, decay=10, thresh=0, bandwidth=lambda x: x.mean(0))
+    G = build_graph(data, decay=10, thresh=0, bandwidth=lambda x: x.mean(1))
     n = G.data.shape[0]
     K = G.build_kernel_to_data(data[:n // 2, :])
     assert(K.shape == (n // 2, n))