Repository: scipy/scipy
File: scipy/spatial/tests/test_kdtree.py (Python)

Code before:

    def test_multiple_radius(self):
        rs = np.exp(np.linspace(np.log(0.01), np.log(10), 3))
        results = self.T1.count_neighbors(self.T2, rs)
        assert_(np.all(np.diff(results) >= 0))
        for r, result in zip(rs, results):
            assert_equal(self.T1.count_neighbors(self.T2, r), result)

@KDTreeTest
class _Test_count_neighbors(count_neighbors_consistency):
    def setup_method(self):
        n = 50
        m = 2
        np.random.seed(1234)
        self.T1 = self.kdtree_type(np.random.randn(n, m), leafsize=2)
        self.T2 = self.kdtree_type(np.random.randn(n, m), leafsize=2)


class sparse_distance_matrix_consistency:

    def distance(self, a, b, p):
        return minkowski_distance(a, b, p)

    def test_consistency_with_neighbors(self):
        M = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dok_array')
        r = self.T1.query_ball_tree(self.T2, self.r)
        for i, l in enumerate(r):
            for j in l:
                assert_almost_equal(
                    M[i, j],
                    self.distance(self.T1.data[i], self.T2.data[j], self.p),
                    decimal=14
                )
        for ((i, j), d) in M.items():
            assert_(j in r[i])

    def test_zero_distance(self):
        # raises an exception for bug 870 (FIXME: Does it?)
        self.T1.sparse_distance_matrix(self.T1, self.r, output_type='dok_array')

    def test_consistency(self):
        # Test consistency with a distance_matrix
        M1 = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dok_array')
        expected = distance_matrix(self.T1.data, self.T2.data)
        expected[expected > self.r] = 0
        assert_array_almost_equal(M1.toarray(), expected, decimal=14)

    def test_against_logic_error_regression(self):
        # regression test for gh-5077 logic error
        np.random.seed(0)
        too_many = np.array(np.random.randn(18, 2), dtype=int)

Review comment (suggestion):
here and elsewhere I did kind of wonder if it would make sense to parametrize over the matrix/array types until the deprecation is over, but maybe that's overkill

Code after:
                np.sum([len(l) for l in self.T1.query_ball_tree(self.T2, r)]))

    def test_multiple_radius(self):
        rs = np.exp(np.linspace(np.log(0.01), np.log(10), 3))
        results = self.T1.count_neighbors(self.T2, rs)
        assert_(np.all(np.diff(results) >= 0))
        for r, result in zip(rs, results):
            assert_equal(self.T1.count_neighbors(self.T2, r), result)

@KDTreeTest
class _Test_count_neighbors(count_neighbors_consistency):
    def setup_method(self):
        n = 50
        m = 2
        np.random.seed(1234)
        self.T1 = self.kdtree_type(np.random.randn(n, m), leafsize=2)
        self.T2 = self.kdtree_type(np.random.randn(n, m), leafsize=2)


class sparse_distance_matrix_consistency:

    def distance(self, a, b, p):
        return minkowski_distance(a, b, p)

    def test_consistency_with_neighbors(self):
        M = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dok_array')
        r = self.T1.query_ball_tree(self.T2, self.r)
        for i, l in enumerate(r):
            for j in l:
                assert_almost_equal(
                    M[i, j],
                    self.distance(self.T1.data[i], self.T2.data[j], self.p),
                    decimal=14
                )
        for ((i, j), d) in M.items():
            assert_(j in r[i])

    def test_zero_distance(self):
        # raises an exception for bug 870 (FIXME: Does it?)
        self.T1.sparse_distance_matrix(self.T1, self.r, output_type='dok_array')

    def test_consistency(self):
        # Test consistency with a distance_matrix
        M1 = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dok_array')
        expected = distance_matrix(self.T1.data, self.T2.data)
        expected[expected > self.r] = 0
        assert_array_almost_equal(M1.toarray(), expected, decimal=14)

    def test_against_logic_error_regression(self):
        # regression test for gh-5077 logic error
        np.random.seed(0)