def test_force2_slow():
from kdcount import force_kernels
from time import time
numpy.random.seed(13)
pos = numpy.random.uniform(size=(32 * 32 * 32 * 8, 3))
tt = KDTree(pos[:1], thresh=1)
tree = KDTree(pos, thresh=1)
mass = KDAttr(tree, numpy.ones(pos.shape[0]))
xmass = KDAttr(tree, pos * mass.input)
print(tt.root, tt.root.min, tt.root.max)
t0 = time()
force2 = tree.root.force2(force_kernels['count'],
tt.root,
mass,
xmass,
1.0 / len(pos)**0.3333 * 8,
eta=1.)
t2 = time() - t0
print('time', t2)
t0 = time()
force1 = tree.root.force(force_kernels['count'],
pos[:1],
mass,
xmass,
1.0 / len(pos)**0.3333 * 8,
eta=1.)
t1 = time() - t0
print('time', t1)
# FIXME: add more tests
assert_array_equal(force1, force2)
def test_force():
from kdcount import force_kernels
pos = numpy.array([[0., 0., 0.], [1., 1., 1.], [2., 2., 2.], [3., 3., 3.]])
#pos = numpy.arange(0, 4, dtype='f8').reshape(-1, 1)
tree = KDTree(pos, thresh=1)
mass = KDAttr(tree, numpy.ones(pos.shape[0]))
xmass = KDAttr(tree, pos * mass.input)
force = tree.root.force(force_kernels['plummer'](1),
pos,
mass,
xmass,
2.,
eta=0.1)
# FIXME: add more tests
print(force)
def test_empty():
pos = numpy.arange(1000).astype('f4').reshape(-1, 1)
shape = ()
data = numpy.ones((len(pos)), dtype=('f4', shape))
tree = KDTree(pos, ind=[])
attr = KDAttr(tree, data)
assert_equal(attr[tree.root], 0)
def test_enum_count_weighted():
numpy.random.seed(1234)
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f4')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f4')
w1 = numpy.ones(len(pos1))
w2 = numpy.ones(len(pos2))
tree1 = KDTree(pos1)
tree2 = KDTree(pos2)
a1 = KDAttr(tree1, w1)
a2 = KDAttr(tree2, w2)
N = [0]
def process1(r, i, j):
N[0] += len(r)
tree1.root.enum(tree2.root, rmax=1.0, process=process1)
c, w = tree1.root.count(tree2.root, r=1.0, attrs=(a1, a2))
assert_equal(N[0], c)
assert_equal(N[0], w)
def test_force_slow():
from kdcount import force_kernels
numpy.random.seed(13)
pos = numpy.random.uniform(size=(32 * 32 * 32, 3))
tree = KDTree(pos, thresh=1)
mass = KDAttr(tree, numpy.ones(pos.shape[0]))
xmass = KDAttr(tree, pos * mass.input)
#force = tree.root.force(force_kernels['plummer'](1), pos, mass, xmass, 1.0 / len(pos) ** 0.3333 * 4, eta=0.1)
force = tree.root.force(force_kernels['count'],
pos,
mass,
xmass,
1.0 / len(pos)**0.3333 * 4,
eta=0.1)
# FIXME: add more tests
print(force)
def test_constattr():
pos = numpy.arange(100).astype('f4').reshape(-1, 1)
shapes = [(), (1, ), (2, )]
for shape in shapes:
data = constant_array((len(pos)), dtype=('f4', shape))
data.value[...] = 1.0
tree = KDTree(pos)
attr = KDAttr(tree, data)
assert_equal(tree.root.index, 0)
assert_equal(tree.root.less.index, 1)
assert_equal(tree.root.greater.index, 2)
assert_equal(attr.buffer.shape[0], tree.size)
assert_equal(attr.buffer.shape[1:], shape)
assert_equal(attr[tree.root], data.sum(axis=0))