Newer
Older
dtype = np.dtype(dtype)
nparticles = 10**5
particles = make_normal_particle_array(actx.queue, nparticles, dims, dtype)
np_particles = actx.to_numpy(particles)
pt.plot(np_particles[0], np_particles[1], "x")
actx.queue.finish()
tree, _ = tb(actx.queue, particles, max_particles_in_box=30, debug=True)
ball_centers = make_normal_particle_array(actx.queue, nballs, dims, dtype)
ball_radii = 0.1 + actx.zeros(nballs, dtype)
from boxtree.area_query import (
LeavesToBallsLookupBuilder, SpaceInvaderQueryBuilder)
siqb = SpaceInvaderQueryBuilder(actx.context)
# We can use leaves-to-balls lookup to get the set of overlapping balls for
# each box, and from there to compute the outer space invader distance.
lblb = LeavesToBallsLookupBuilder(actx.context)
siq, _ = siqb(actx.queue, tree, ball_centers, ball_radii)
lbl, _ = lblb(actx.queue, tree, ball_centers, ball_radii)
tree = tree.get(queue=actx.queue)
siq = siq.get(queue=actx.queue)
lbl = lbl.get(queue=actx.queue)
ball_centers = np.array([actx.to_numpy(x) for x in ball_centers])
ball_radii = actx.to_numpy(ball_radii)
# Find leaf boxes.
from boxtree import box_flags_enum
outer_space_invader_dist = np.zeros(tree.nboxes)
for ibox in range(tree.nboxes):
# We only want leaves here.
if tree.box_flags[ibox] & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES:
continue
start, end = lbl.balls_near_box_starts[ibox:ibox + 2]
space_invaders = lbl.balls_near_box_lists[start:end]
if len(space_invaders) > 0:
outer_space_invader_dist[ibox] = np.max(np.abs(
tree.box_centers[:, ibox].reshape((-1, 1))
- ball_centers[:, space_invaders]))
assert np.allclose(siq, outer_space_invader_dist)
# }}}
# {{{ test_same_tree_with_zero_weight_particles
@pytest.mark.parametrize("dims", [2, 3])
def test_same_tree_with_zero_weight_particles(actx_factory, dims):
actx = actx_factory()
ntargets_values = [300, 400, 500]
stick_out_factors = [0, 0.1, 0.3, 1]
nsources = 20
from boxtree import TreeBuilder
for stick_out_factor in stick_out_factors:
for ntargets in [40]:
sources = rng.random((dims, nsources))**2
sources[:, 0] = -0.1
sources[:, 1] = 1.1
targets = rng.random((dims, max(ntargets_values)))[:, :ntargets].copy()
target_radii = rng.random(max(ntargets_values))[:ntargets]
sources = actx.from_numpy(sources)
targets = actx.from_numpy(targets)
refine_weights = actx.empty(nsources + ntargets, np.int32)
refine_weights[:nsources] = 1
refine_weights[nsources:] = 0
tree, _ = tb(actx.queue, sources, targets=targets,
target_radii=target_radii,
stick_out_factor=stick_out_factor,
max_leaf_refine_weight=10,
refine_weights=refine_weights,
debug=True)
trees.append(tree)
print("TREE:", tree.nboxes)
if 0:
import matplotlib.pyplot as plt
for tree in trees:
plt.figure()
tree.plot()
plt.show()
# }}}
def test_max_levels_error(actx_factory):
actx = actx_factory()
sources = [actx.zeros(11, np.float64) for i in range(2)]
from boxtree.tree_build import MaxLevelsExceeded
with pytest.raises(MaxLevelsExceeded):
_tree, _ = tb(actx.queue, sources, max_particles_in_box=10, debug=True)
# $ python test_tree.py 'test_routine(cl.create_some_context)'
if __name__ == "__main__":
if len(sys.argv) > 1:
exec(sys.argv[1])
else: