TST Add test to make sure BallTree remains picklable even with callable metric

raghavrv · raghavrv · commit b42a14813561 · 2015-05-03T23:52:13.000+05:30
diff --git a/sklearn/neighbors/tests/test_ball_tree.py b/sklearn/neighbors/tests/test_ball_tree.py
@@ -1,3 +1,4 @@
+import pickle
 import numpy as np
 from numpy.testing import assert_array_almost_equal
 from sklearn.neighbors.ball_tree import (BallTree, NeighborsHeap,
@@ -29,6 +30,10 @@
                    'sokalsneath']
 
 
+def dist_func(x1, x2, p):
+    return np.sum((x1 - x2) ** p) ** (1. / p)
+
+
 def brute_force_neighbors(X, Y, k, metric, **kwargs):
     D = DistanceMetric.get_metric(metric, **kwargs).pairwise(Y, X)
     ind = np.argsort(D, axis=1)[:, :k]
@@ -216,19 +221,32 @@ def check_two_point(r, dualtree):
 
 
 def test_ball_tree_pickle():
-    import pickle
     np.random.seed(0)
     X = np.random.random((10, 3))
+
     bt1 = BallTree(X, leaf_size=1)
+    # Test if BallTree with callable metric is picklable
+    bt1_pyfunc = BallTree(X, metric=dist_func, leaf_size=1, p=2)
+
     ind1, dist1 = bt1.query(X)
+    ind1_pyfunc, dist1_pyfunc = bt1_pyfunc.query(X)
 
     def check_pickle_protocol(protocol):
         s = pickle.dumps(bt1, protocol=protocol)
         bt2 = pickle.loads(s)
+
+        s_pyfunc = pickle.dumps(bt1_pyfunc, protocol=protocol)
+        bt2_pyfunc = pickle.loads(s_pyfunc)
+
         ind2, dist2 = bt2.query(X)
+        ind2_pyfunc, dist2_pyfunc = bt2_pyfunc.query(X)
+
         assert_array_almost_equal(ind1, ind2)
         assert_array_almost_equal(dist1, dist2)
 
+        assert_array_almost_equal(ind1_pyfunc, ind2_pyfunc)
+        assert_array_almost_equal(dist1_pyfunc, dist2_pyfunc)
+
     for protocol in (0, 1, 2):
         yield check_pickle_protocol, protocol
 
