test_parallel_utils.py 文件源码

def test_kdtree_parallel_distribute(ndim=2, npts=50):
    total_npts = npts
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()
    if rank == 0:
        total_pts = np.random.rand(total_npts, ndim).astype('float64')
    else:
        total_pts = None
    pts, idx, le, re, ple, pre = parallel_utils.py_kdtree_parallel_distribute(total_pts)
    total_pts = comm.bcast(total_pts, root=0)
    assert_equal(pts.shape[0], idx.size)
    np.testing.assert_array_equal(pts, total_pts[idx,:])
    for d in range(ndim):
        assert_less_equal(pts[:,d], re[d])
        assert_less_equal(le[d], pts[:,d])