test_neighbors.py 文件源码

def test_same_knn_parallel():
    X, y = datasets.make_classification(n_samples=30, n_features=5,
                                        n_redundant=0, random_state=0)
    X_train, X_test, y_train, y_test = train_test_split(X, y)

    def check_same_knn_parallel(algorithm):
        clf = neighbors.KNeighborsClassifier(n_neighbors=3,
                                             algorithm=algorithm)
        clf.fit(X_train, y_train)
        y = clf.predict(X_test)
        dist, ind = clf.kneighbors(X_test)
        graph = clf.kneighbors_graph(X_test, mode='distance').toarray()

        clf.set_params(n_jobs=3)
        clf.fit(X_train, y_train)
        y_parallel = clf.predict(X_test)
        dist_parallel, ind_parallel = clf.kneighbors(X_test)
        graph_parallel = \
            clf.kneighbors_graph(X_test, mode='distance').toarray()

        assert_array_equal(y, y_parallel)
        assert_array_almost_equal(dist, dist_parallel)
        assert_array_equal(ind, ind_parallel)
        assert_array_almost_equal(graph, graph_parallel)

    for algorithm in ALGORITHMS:
        yield check_same_knn_parallel, algorithm