def test_area_query_elwise(ctx_getter, dims, do_plot=False):
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nparticles = 10**5
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
from boxtree.area_query import (
AreaQueryElementwiseTemplate, PeerListFinder)
template = AreaQueryElementwiseTemplate(
extra_args="""
coord_t *ball_radii,
%for ax in AXIS_NAMES[:dimensions]:
coord_t *ball_${ax},
%endfor
""",
ball_center_and_radius_expr="""
%for ax in AXIS_NAMES[:dimensions]:
${ball_center}.${ax} = ball_${ax}[${i}];
%endfor
${ball_radius} = ball_radii[${i}];
""",
leaf_found_op="")
peer_lists, evt = PeerListFinder(ctx)(queue, tree)
kernel = template.generate(
ctx,
dims,
tree.coord_dtype,
tree.box_id_dtype,
peer_lists.peer_list_starts.dtype,
tree.nlevels)
evt = kernel(
*template.unwrap_args(
tree, peer_lists, ball_radii, *ball_centers),
queue=queue,
wait_for=[evt],
range=slice(len(ball_radii)))
cl.wait_for_events([evt])
def test_leaves_to_balls_query(ctx_factory, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nparticles = 10**5
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
from boxtree.area_query import LeavesToBallsLookupBuilder
lblb = LeavesToBallsLookupBuilder(ctx)
lbl, _ = lblb(queue, tree, ball_centers, ball_radii)
# get data to host for test
tree = tree.get(queue=queue)
lbl = lbl.get(queue=queue)
ball_centers = np.array([x.get() for x in ball_centers]).T
ball_radii = ball_radii.get()
assert len(lbl.balls_near_box_starts) == tree.nboxes + 1
from boxtree import box_flags_enum
for ibox in range(tree.nboxes):
# We only want leaves here.
if tree.box_flags[ibox] & box_flags_enum.HAS_CHILDREN:
continue
box_center = tree.box_centers[:, ibox]
ext_l, ext_h = tree.get_box_extent(ibox)
box_rad = 0.5 * (ext_h - ext_l)[0]
linf_circle_dists = np.max(np.abs(ball_centers - box_center), axis=-1)
near_circles, = np.where(linf_circle_dists - ball_radii < box_rad)
start, end = lbl.balls_near_box_starts[ibox:ibox + 2]
assert sorted(
lbl.balls_near_box_lists[start:end]) == sorted(near_circles)
def test_leaves_to_balls_query(ctx_getter, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nparticles = 10**5
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
from boxtree.area_query import LeavesToBallsLookupBuilder
lblb = LeavesToBallsLookupBuilder(ctx)
lbl, _ = lblb(queue, tree, ball_centers, ball_radii)
# get data to host for test
tree = tree.get(queue=queue)
lbl = lbl.get(queue=queue)
ball_centers = np.array([x.get() for x in ball_centers]).T
ball_radii = ball_radii.get()
assert len(lbl.balls_near_box_starts) == tree.nboxes + 1
from boxtree import box_flags_enum
for ibox in range(tree.nboxes):
# We only want leaves here.
if tree.box_flags[ibox] & box_flags_enum.HAS_CHILDREN:
continue
box_center = tree.box_centers[:, ibox]
ext_l, ext_h = tree.get_box_extent(ibox)
box_rad = 0.5*(ext_h-ext_l)[0]
linf_circle_dists = np.max(np.abs(ball_centers-box_center), axis=-1)
near_circles, = np.where(linf_circle_dists - ball_radii < box_rad)
start, end = lbl.balls_near_box_starts[ibox:ibox+2]
assert sorted(lbl.balls_near_box_lists[start:end]) == sorted(near_circles)
def test_area_query_elwise(ctx_factory, dims, do_plot=False):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nparticles = 10**5
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
from boxtree.area_query import (AreaQueryElementwiseTemplate,
PeerListFinder)
template = AreaQueryElementwiseTemplate(extra_args="""
coord_t *ball_radii,
%for ax in AXIS_NAMES[:dimensions]:
coord_t *ball_${ax},
%endfor
""",
ball_center_and_radius_expr="""
%for ax in AXIS_NAMES[:dimensions]:
${ball_center}.${ax} = ball_${ax}[${i}];
%endfor
${ball_radius} = ball_radii[${i}];
""",
leaf_found_op="")
peer_lists, evt = PeerListFinder(ctx)(queue, tree)
kernel = template.generate(ctx, dims, tree.coord_dtype, tree.box_id_dtype,
peer_lists.peer_list_starts.dtype, tree.nlevels)
evt = kernel(*template.unwrap_args(tree, peer_lists, ball_radii,
*ball_centers),
queue=queue,
wait_for=[evt],
range=slice(len(ball_radii)))
cl.wait_for_events([evt])
def test_bounding_box(ctx_getter, dtype, dims, nparticles):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
from boxtree.tools import AXIS_NAMES
from boxtree.bounding_box import BoundingBoxFinder
bbf = BoundingBoxFinder(ctx)
axis_names = AXIS_NAMES[:dims]
logger.info("%s - %s %s" % (dtype, dims, nparticles))
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
bbox_min = [np.min(x.get()) for x in particles]
bbox_max = [np.max(x.get()) for x in particles]
bbox_cl, evt = bbf(particles, radii=None)
bbox_cl = bbox_cl.get()
bbox_min_cl = np.empty(dims, dtype)
bbox_max_cl = np.empty(dims, dtype)
for i, ax in enumerate(axis_names):
bbox_min_cl[i] = bbox_cl["min_"+ax]
bbox_max_cl[i] = bbox_cl["max_"+ax]
assert (bbox_min == bbox_min_cl).all()
assert (bbox_max == bbox_max_cl).all()
def test_bounding_box(ctx_factory, dtype, dims, nparticles):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
from boxtree.tools import AXIS_NAMES
from boxtree.bounding_box import BoundingBoxFinder
bbf = BoundingBoxFinder(ctx)
axis_names = AXIS_NAMES[:dims]
logger.info("%s - %s %s" % (dtype, dims, nparticles))
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
bbox_min = [np.min(x.get()) for x in particles]
bbox_max = [np.max(x.get()) for x in particles]
bbox_cl, evt = bbf(particles, radii=None)
bbox_cl = bbox_cl.get()
bbox_min_cl = np.empty(dims, dtype)
bbox_max_cl = np.empty(dims, dtype)
for i, ax in enumerate(axis_names):
bbox_min_cl[i] = bbox_cl["min_" + ax]
bbox_max_cl[i] = bbox_cl["max_" + ax]
assert (bbox_min == bbox_min_cl).all()
assert (bbox_max == bbox_max_cl).all()
def test_area_query_balls_outside_bbox(ctx_getter, dims, do_plot=False):
"""
The input to the area query includes balls whose centers are not within
the tree bounding box.
"""
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nparticles = 10**4
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx, seed=13)
bbox_min = tree.bounding_box[0].min()
bbox_max = tree.bounding_box[1].max()
from pytools.obj_array import make_obj_array
ball_centers = make_obj_array([
rng.uniform(queue, nballs, dtype=dtype, a=bbox_min-1, b=bbox_max+1)
for i in range(dims)])
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
run_area_query_test(ctx, queue, tree, ball_centers, ball_radii)
def test_area_query(ctx_factory, dims, do_plot=False):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nparticles = 10**5
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
run_area_query_test(ctx, queue, tree, ball_centers, ball_radii)
def test_area_query(ctx_getter, dims, do_plot=False):
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nparticles = 10**5
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
run_area_query_test(ctx, queue, tree, ball_centers, ball_radii)
def test_area_query_balls_outside_bbox(ctx_factory, dims, do_plot=False):
"""
The input to the area query includes balls whose centers are not within
the tree bounding box.
"""
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nparticles = 10**4
dtype = np.float64
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx, seed=13)
bbox_min = tree.bounding_box[0].min()
bbox_max = tree.bounding_box[1].max()
from pytools.obj_array import make_obj_array
ball_centers = make_obj_array([
rng.uniform(queue, nballs, dtype=dtype, a=bbox_min - 1, b=bbox_max + 1)
for i in range(dims)
])
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
run_area_query_test(ctx, queue, tree, ball_centers, ball_radii)
def run_build_test(builder, queue, dims, dtype, nparticles, do_plot,
max_particles_in_box=None, max_leaf_refine_weight=None,
refine_weights=None, **kwargs):
dtype = np.dtype(dtype)
if dtype == np.float32:
tol = 1e-4
elif dtype == np.float64:
tol = 1e-12
else:
raise RuntimeError("unsupported dtype: %s" % dtype)
logger.info(75*"-")
if max_particles_in_box is not None:
logger.info("%dD %s - %d particles - max %d per box - %s" % (
dims, dtype.type.__name__, nparticles, max_particles_in_box,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
else:
logger.info("%dD %s - %d particles - max leaf weight %d - %s" % (
dims, dtype.type.__name__, nparticles, max_leaf_refine_weight,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
logger.info(75*"-")
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
queue.finish()
tree, _ = builder(queue, particles,
max_particles_in_box=max_particles_in_box,
refine_weights=refine_weights,
max_leaf_refine_weight=max_leaf_refine_weight,
debug=True, **kwargs)
tree = tree.get(queue=queue)
sorted_particles = np.array(list(tree.sources))
unsorted_particles = np.array([pi.get() for pi in particles])
assert (sorted_particles
== unsorted_particles[:, tree.user_source_ids]).all()
if refine_weights is not None:
refine_weights_reordered = refine_weights.get()[tree.user_source_ids]
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
from boxtree import box_flags_enum as bfe
scaled_tol = tol*tree.root_extent
for ibox in range(tree.nboxes):
# Empty boxes exist in non-pruned trees--which themselves are undocumented.
# These boxes will fail these tests.
if not (tree.box_flags[ibox] & bfe.HAS_OWN_SRCNTGTS):
continue
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >= tree.bounding_box[0] - scaled_tol).all(), (
ibox, extent_low, tree.bounding_box[0])
assert (extent_high <= tree.bounding_box[1] + scaled_tol).all(), (
ibox, extent_high, tree.bounding_box[1])
start = tree.box_source_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
box_particles = sorted_particles[:,
start:start+tree.box_source_counts_cumul[ibox]]
good = (
(box_particles < extent_high[:, np.newaxis] + scaled_tol)
& (extent_low[:, np.newaxis] - scaled_tol <= box_particles))
all_good_here = good.all()
if do_plot and not all_good_here and all_good_so_far:
pt.plot(
box_particles[0, np.where(~good)[1]],
box_particles[1, np.where(~good)[1]], "ro")
plotter.draw_box(ibox, edgecolor="red")
if not all_good_here:
print("BAD BOX", ibox)
if not (tree.box_flags[ibox] & bfe.HAS_CHILDREN):
# Check that leaf particle density is as promised.
nparticles_in_box = tree.box_source_counts_cumul[ibox]
if max_particles_in_box is not None:
if nparticles_in_box > max_particles_in_box:
print("too many particles ({0} > {1}); box {2}".format(
nparticles_in_box, max_particles_in_box, ibox))
all_good_here = False
else:
assert refine_weights is not None
box_weight = np.sum(
refine_weights_reordered[start:start+nparticles_in_box])
if box_weight > max_leaf_refine_weight:
print("refine weight exceeded ({0} > {1}); box {2}".format(
box_weight, max_leaf_refine_weight, ibox))
all_good_here = False
all_good_so_far = all_good_so_far and all_good_here
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
assert all_good_so_far
def test_source_target_tree(ctx_factory, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nsources = 2 * 10**5
ntargets = 3 * 10**5
dtype = np.float64
sources = make_normal_particle_array(queue, nsources, dims, dtype, seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype, seed=19)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
pt.show()
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue,
sources,
targets=targets,
max_particles_in_box=10,
debug=True)
tree = tree.get(queue=queue)
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
assert (sorted_sources == unsorted_sources[:, tree.user_source_ids]).all()
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets,
dtype=np.intp)
assert (sorted_targets == unsorted_targets[:, user_target_ids]).all()
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
tol = 1e-15
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >=
tree.bounding_box[0] - 1e-12 * tree.root_extent).all(), ibox
assert (extent_high <=
tree.bounding_box[1] + 1e-12 * tree.root_extent).all(), ibox
src_start = tree.box_source_starts[ibox]
tgt_start = tree.box_target_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox] +
np.sum(tree.box_source_counts_cumul[existing_children]) ==
tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox] +
np.sum(tree.box_target_counts_cumul[existing_children]) ==
tree.box_target_counts_cumul[ibox])
for what, particles in [
("sources", sorted_sources[:, src_start:src_start +
tree.box_source_counts_cumul[ibox]]),
("targets", sorted_targets[:, tgt_start:tgt_start +
tree.box_target_counts_cumul[ibox]]),
]:
good = ((particles < extent_high[:, np.newaxis] + tol)
&
(extent_low[:, np.newaxis] - tol <= particles)).all(axis=0)
all_good_here = good.all()
if do_plot and not all_good_here:
pt.plot(particles[0, np.where(~good)[0]],
particles[1, np.where(~good)[0]], "ro")
plotter.draw_box(ibox, edgecolor="red")
pt.show()
if not all_good_here:
print("BAD BOX %s %d" % (what, ibox))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_so_far
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
def test_space_invader_query(ctx_factory, dims, dtype, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
dtype = np.dtype(dtype)
nparticles = 10**5
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
from boxtree.area_query import (LeavesToBallsLookupBuilder,
SpaceInvaderQueryBuilder)
siqb = SpaceInvaderQueryBuilder(ctx)
# 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(ctx)
siq, _ = siqb(queue, tree, ball_centers, ball_radii)
lbl, _ = lblb(queue, tree, ball_centers, ball_radii)
# get data to host for test
tree = tree.get(queue=queue)
siq = siq.get(queue=queue)
lbl = lbl.get(queue=queue)
ball_centers = np.array([x.get() for x in ball_centers])
ball_radii = ball_radii.get()
# 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_CHILDREN:
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)
def test_extent_tree(ctx_factory, dims, extent_norm, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nsources = 100000
ntargets = 200000
dtype = np.float64
npoint_sources_per_source = 16
sources = make_normal_particle_array(queue, nsources, dims, dtype, seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype, seed=19)
refine_weights = cl.array.zeros(queue, nsources + ntargets, np.int32)
refine_weights[:nsources] = 1
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=13)
source_radii = 2**rng.uniform(queue, nsources, dtype=dtype, a=-10, b=0)
target_radii = 2**rng.uniform(queue, ntargets, dtype=dtype, a=-10, b=0)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
dev_tree, _ = tb(
queue,
sources,
targets=targets,
source_radii=source_radii,
target_radii=target_radii,
extent_norm=extent_norm,
refine_weights=refine_weights,
max_leaf_refine_weight=20,
#max_particles_in_box=10,
# Set artificially small, to exercise the reallocation code.
nboxes_guess=10,
debug=True,
stick_out_factor=0)
logger.info("transfer tree, check orderings")
tree = dev_tree.get(queue=queue)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.gca().set_aspect("equal", "datalim")
pt.show()
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
sorted_source_radii = tree.source_radii
sorted_target_radii = tree.target_radii
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
unsorted_source_radii = source_radii.get()
unsorted_target_radii = target_radii.get()
assert (sorted_sources == unsorted_sources[:, tree.user_source_ids]).all()
assert (sorted_source_radii == unsorted_source_radii[tree.user_source_ids]
).all()
# {{{ test box structure, stick-out criterion
logger.info("test box structure, stick-out criterion")
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets,
dtype=np.intp)
if ntargets:
assert (sorted_targets == unsorted_targets[:, user_target_ids]).all()
assert (sorted_target_radii == unsorted_target_radii[user_target_ids]
).all()
all_good_so_far = True
# {{{ check sources, targets
assert np.sum(tree.box_source_counts_nonchild) == nsources
assert np.sum(tree.box_target_counts_nonchild) == ntargets
for ibox in range(tree.nboxes):
kid_sum = sum(tree.box_target_counts_cumul[ichild_box]
for ichild_box in tree.box_child_ids[:, ibox]
if ichild_box != 0)
assert (tree.box_target_counts_cumul[ibox] == (
tree.box_target_counts_nonchild[ibox] + kid_sum)), ibox
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >=
tree.bounding_box[0] - 1e-12 * tree.root_extent).all(), ibox
assert (extent_high <=
tree.bounding_box[1] + 1e-12 * tree.root_extent).all(), ibox
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox] +
np.sum(tree.box_source_counts_cumul[existing_children]) ==
tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox] +
np.sum(tree.box_target_counts_cumul[existing_children]) ==
tree.box_target_counts_cumul[ibox])
del existing_children
del box_children
for ibox in range(tree.nboxes):
lev = int(tree.box_levels[ibox])
box_radius = 0.5 * tree.root_extent / (1 << lev)
box_center = tree.box_centers[:, ibox]
extent_low = box_center - box_radius
extent_high = box_center + box_radius
stick_out_dist = tree.stick_out_factor * box_radius
radius_with_stickout = (1 + tree.stick_out_factor) * box_radius
for what, starts, counts, points, radii in [
("source", tree.box_source_starts, tree.box_source_counts_cumul,
sorted_sources, sorted_source_radii),
("target", tree.box_target_starts, tree.box_target_counts_cumul,
sorted_targets, sorted_target_radii),
]:
bstart = starts[ibox]
bslice = slice(bstart, bstart + counts[ibox])
check_particles = points[:, bslice]
check_radii = radii[bslice]
if extent_norm == "linf":
good = ((check_particles + check_radii <
extent_high[:, np.newaxis] + stick_out_dist)
& # noqa: W504
(extent_low[:, np.newaxis] - stick_out_dist <=
check_particles - check_radii)).all(axis=0)
elif extent_norm == "l2":
center_dists = np.sqrt(
np.sum((check_particles - box_center.reshape(-1, 1))**2,
axis=0))
good = ((center_dists + check_radii)**2 <
dims * radius_with_stickout**2)
else:
raise ValueError("unexpected value of extent_norm")
all_good_here = good.all()
if not all_good_here:
print("BAD BOX %s %d level %d" %
(what, ibox, tree.box_levels[ibox]))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_here
# }}}
assert all_good_so_far
# }}}
# {{{ create, link point sources
logger.info("creating point sources")
np.random.seed(20)
from pytools.obj_array import make_obj_array
point_sources = make_obj_array([
cl.array.to_device(
queue, unsorted_sources[i][:, np.newaxis] +
unsorted_source_radii[:, np.newaxis] * np.random.uniform(
-1, 1, size=(nsources, npoint_sources_per_source)))
for i in range(dims)
])
point_source_starts = cl.array.arange(queue,
0, (nsources + 1) *
npoint_sources_per_source,
npoint_sources_per_source,
dtype=tree.particle_id_dtype)
from boxtree.tree import link_point_sources
dev_tree = link_point_sources(queue,
dev_tree,
point_source_starts,
point_sources,
debug=True)
def run_build_test(builder,
queue,
dims,
dtype,
nparticles,
do_plot,
max_particles_in_box=None,
max_leaf_refine_weight=None,
refine_weights=None,
**kwargs):
dtype = np.dtype(dtype)
if dtype == np.float32:
tol = 1e-4
elif dtype == np.float64:
tol = 1e-12
else:
raise RuntimeError("unsupported dtype: %s" % dtype)
logger.info(75 * "-")
if max_particles_in_box is not None:
logger.info(
"%dD %s - %d particles - max %d per box - %s" %
(dims, dtype.type.__name__, nparticles, max_particles_in_box,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
else:
logger.info(
"%dD %s - %d particles - max leaf weight %d - %s" %
(dims, dtype.type.__name__, nparticles, max_leaf_refine_weight,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
logger.info(75 * "-")
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
queue.finish()
tree, _ = builder(queue,
particles,
max_particles_in_box=max_particles_in_box,
refine_weights=refine_weights,
max_leaf_refine_weight=max_leaf_refine_weight,
debug=True,
**kwargs)
tree = tree.get(queue=queue)
sorted_particles = np.array(list(tree.sources))
unsorted_particles = np.array([pi.get() for pi in particles])
assert (
sorted_particles == unsorted_particles[:, tree.user_source_ids]).all()
if refine_weights is not None:
refine_weights_reordered = refine_weights.get()[tree.user_source_ids]
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
from boxtree import box_flags_enum as bfe
scaled_tol = tol * tree.root_extent
for ibox in range(tree.nboxes):
# Empty boxes exist in non-pruned trees--which themselves are undocumented.
# These boxes will fail these tests.
if not (tree.box_flags[ibox] & bfe.HAS_OWN_SRCNTGTS):
continue
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >= tree.bounding_box[0] - scaled_tol).all(), (
ibox, extent_low, tree.bounding_box[0])
assert (extent_high <= tree.bounding_box[1] + scaled_tol).all(), (
ibox, extent_high, tree.bounding_box[1])
start = tree.box_source_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox] +
np.sum(tree.box_source_counts_cumul[existing_children]) ==
tree.box_source_counts_cumul[ibox])
box_particles = sorted_particles[:, start:start +
tree.box_source_counts_cumul[ibox]]
good = ((box_particles < extent_high[:, np.newaxis] + scaled_tol)
& (extent_low[:, np.newaxis] - scaled_tol <= box_particles))
all_good_here = good.all()
if do_plot and not all_good_here and all_good_so_far:
pt.plot(box_particles[0, np.where(~good)[1]],
box_particles[1, np.where(~good)[1]], "ro")
plotter.draw_box(ibox, edgecolor="red")
if not all_good_here:
print("BAD BOX", ibox)
if not (tree.box_flags[ibox] & bfe.HAS_CHILDREN):
# Check that leaf particle density is as promised.
nparticles_in_box = tree.box_source_counts_cumul[ibox]
if max_particles_in_box is not None:
if nparticles_in_box > max_particles_in_box:
print("too many particles ({0} > {1}); box {2}".format(
nparticles_in_box, max_particles_in_box, ibox))
all_good_here = False
else:
assert refine_weights is not None
box_weight = np.sum(
refine_weights_reordered[start:start + nparticles_in_box])
if box_weight > max_leaf_refine_weight:
print("refine weight exceeded ({0} > {1}); box {2}".format(
box_weight, max_leaf_refine_weight, ibox))
all_good_here = False
all_good_so_far = all_good_so_far and all_good_here
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
assert all_good_so_far
def test_extent_tree(ctx_getter, dims, extent_norm, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 100000
ntargets = 200000
dtype = np.float64
npoint_sources_per_source = 16
sources = make_normal_particle_array(queue, nsources, dims, dtype,
seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype,
seed=19)
refine_weights = cl.array.zeros(queue, nsources+ntargets, np.int32)
refine_weights[:nsources] = 1
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=13)
source_radii = 2**rng.uniform(queue, nsources, dtype=dtype,
a=-10, b=0)
target_radii = 2**rng.uniform(queue, ntargets, dtype=dtype,
a=-10, b=0)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
dev_tree, _ = tb(queue, sources, targets=targets,
source_radii=source_radii,
target_radii=target_radii,
extent_norm=extent_norm,
refine_weights=refine_weights,
max_leaf_refine_weight=20,
#max_particles_in_box=10,
# Set artificially small, to exercise the reallocation code.
nboxes_guess=10,
debug=True,
stick_out_factor=0)
logger.info("transfer tree, check orderings")
tree = dev_tree.get(queue=queue)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.gca().set_aspect("equal", "datalim")
pt.show()
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
sorted_source_radii = tree.source_radii
sorted_target_radii = tree.target_radii
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
unsorted_source_radii = source_radii.get()
unsorted_target_radii = target_radii.get()
assert (sorted_sources
== unsorted_sources[:, tree.user_source_ids]).all()
assert (sorted_source_radii
== unsorted_source_radii[tree.user_source_ids]).all()
# {{{ test box structure, stick-out criterion
logger.info("test box structure, stick-out criterion")
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets, dtype=np.intp)
if ntargets:
assert (sorted_targets
== unsorted_targets[:, user_target_ids]).all()
assert (sorted_target_radii
== unsorted_target_radii[user_target_ids]).all()
all_good_so_far = True
# {{{ check sources, targets
assert np.sum(tree.box_source_counts_nonchild) == nsources
assert np.sum(tree.box_target_counts_nonchild) == ntargets
for ibox in range(tree.nboxes):
kid_sum = sum(
tree.box_target_counts_cumul[ichild_box]
for ichild_box in tree.box_child_ids[:, ibox]
if ichild_box != 0)
assert (
tree.box_target_counts_cumul[ibox]
== (
tree.box_target_counts_nonchild[ibox]
+ kid_sum)), ibox
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low
>= tree.bounding_box[0] - 1e-12*tree.root_extent).all(), ibox
assert (extent_high
<= tree.bounding_box[1] + 1e-12*tree.root_extent).all(), ibox
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox]
+ np.sum(tree.box_target_counts_cumul[existing_children])
== tree.box_target_counts_cumul[ibox])
del existing_children
del box_children
for ibox in range(tree.nboxes):
lev = int(tree.box_levels[ibox])
box_radius = 0.5 * tree.root_extent / (1 << lev)
box_center = tree.box_centers[:, ibox]
extent_low = box_center - box_radius
extent_high = box_center + box_radius
stick_out_dist = tree.stick_out_factor * box_radius
radius_with_stickout = (1 + tree.stick_out_factor) * box_radius
for what, starts, counts, points, radii in [
("source", tree.box_source_starts, tree.box_source_counts_cumul,
sorted_sources, sorted_source_radii),
("target", tree.box_target_starts, tree.box_target_counts_cumul,
sorted_targets, sorted_target_radii),
]:
bstart = starts[ibox]
bslice = slice(bstart, bstart+counts[ibox])
check_particles = points[:, bslice]
check_radii = radii[bslice]
if extent_norm == "linf":
good = (
(check_particles + check_radii
< extent_high[:, np.newaxis] + stick_out_dist)
& # noqa: W504
(extent_low[:, np.newaxis] - stick_out_dist
<= check_particles - check_radii)
).all(axis=0)
elif extent_norm == "l2":
center_dists = np.sqrt(
np.sum(
(check_particles - box_center.reshape(-1, 1))**2,
axis=0))
good = (
(center_dists + check_radii)**2
< dims * radius_with_stickout**2)
else:
raise ValueError("unexpected value of extent_norm")
all_good_here = good.all()
if not all_good_here:
print("BAD BOX %s %d level %d"
% (what, ibox, tree.box_levels[ibox]))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_here
# }}}
assert all_good_so_far
# }}}
# {{{ create, link point sources
logger.info("creating point sources")
np.random.seed(20)
from pytools.obj_array import make_obj_array
point_sources = make_obj_array([
cl.array.to_device(queue,
unsorted_sources[i][:, np.newaxis]
+ unsorted_source_radii[:, np.newaxis]
* np.random.uniform(
-1, 1, size=(nsources, npoint_sources_per_source))
)
for i in range(dims)])
point_source_starts = cl.array.arange(queue,
0, (nsources+1)*npoint_sources_per_source, npoint_sources_per_source,
dtype=tree.particle_id_dtype)
from boxtree.tree import link_point_sources
dev_tree = link_point_sources(queue, dev_tree,
point_source_starts, point_sources,
debug=True)
def run_build_test(builder, queue, dims, dtype, nparticles, do_plot,
max_particles_in_box=30, **kwargs):
dtype = np.dtype(dtype)
if dtype == np.float32:
tol = 1e-4
elif dtype == np.float64:
tol = 1e-12
else:
raise RuntimeError("unsupported dtype: %s" % dtype)
if (dtype == np.float32
and dims == 2
and queue.device.platform.name == "Portable Computing Language"):
pytest.xfail("2D float doesn't work on POCL")
logger.info(75*"-")
logger.info("%dD %s - %d particles - max %d per box - %s" % (
dims, dtype.type.__name__, nparticles, max_particles_in_box,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
logger.info(75*"-")
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
queue.finish()
tree, _ = builder(queue, particles,
max_particles_in_box=max_particles_in_box, debug=True,
**kwargs)
tree = tree.get(queue=queue)
sorted_particles = np.array(list(tree.sources))
unsorted_particles = np.array([pi.get() for pi in particles])
assert (sorted_particles
== unsorted_particles[:, tree.user_source_ids]).all()
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
from boxtree import box_flags_enum as bfe
scaled_tol = tol*tree.root_extent
for ibox in range(tree.nboxes):
# Empty boxes exist in non-pruned trees--which themselves are undocumented.
# These boxes will fail these tests.
if not (tree.box_flags[ibox] & bfe.HAS_OWN_SRCNTGTS):
continue
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >= tree.bounding_box[0] - scaled_tol).all(), (
ibox, extent_low, tree.bounding_box[0])
assert (extent_high <= tree.bounding_box[1] + scaled_tol).all(), (
ibox, extent_high, tree.bounding_box[1])
start = tree.box_source_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
box_particles = sorted_particles[:,
start:start+tree.box_source_counts_cumul[ibox]]
good = (
(box_particles < extent_high[:, np.newaxis] + scaled_tol)
&
(extent_low[:, np.newaxis] - scaled_tol <= box_particles)
)
all_good_here = good.all()
if do_plot and not all_good_here and all_good_so_far:
pt.plot(
box_particles[0, np.where(~good)[1]],
box_particles[1, np.where(~good)[1]], "ro")
plotter.draw_box(ibox, edgecolor="red")
if not all_good_here:
print("BAD BOX", ibox)
all_good_so_far = all_good_so_far and all_good_here
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
assert all_good_so_far
def test_tree_connectivity(ctx_getter, dims, sources_are_targets):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
dtype = np.float64
sources = make_normal_particle_array(queue, 1 * 10**5, dims, dtype)
if sources_are_targets:
targets = None
else:
targets = make_normal_particle_array(queue, 2 * 10**5, dims, dtype)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue,
sources,
max_particles_in_box=30,
targets=targets,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree, debug=True)
tree = tree.get(queue=queue)
trav = trav.get(queue=queue)
levels = tree.box_levels
parents = tree.box_parent_ids.T
children = tree.box_child_ids.T
centers = tree.box_centers.T
# {{{ parent and child relations, levels match up
for ibox in range(1, tree.nboxes):
# /!\ Not testing box 0, has no parents
parent = parents[ibox]
assert levels[parent] + 1 == levels[ibox]
assert ibox in children[parent], ibox
# }}}
if 0:
import matplotlib.pyplot as pt
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.show()
# {{{ neighbor_source_boxes (list 1) consists of source boxes
for itgt_box, ibox in enumerate(trav.target_boxes):
start, end = trav.neighbor_source_boxes_starts[itgt_box:itgt_box + 2]
nbl = trav.neighbor_source_boxes_lists[start:end]
if sources_are_targets:
assert ibox in nbl
for jbox in nbl:
assert (0 == children[jbox]).all(), (ibox, jbox, children[jbox])
logger.info("list 1 consists of source boxes")
# }}}
# {{{ separated siblings (list 2) are actually separated
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_siblings_starts[itgt_box:itgt_box + 2]
seps = trav.sep_siblings_lists[start:end]
assert (levels[seps] == levels[tgt_ibox]).all()
# three-ish box radii (half of size)
mindist = 2.5 * 0.5 * 2**-int(levels[tgt_ibox]) * tree.root_extent
icenter = centers[tgt_ibox]
for jbox in seps:
dist = la.norm(centers[jbox] - icenter)
assert dist > mindist, (dist, mindist)
logger.info("separated siblings (list 2) are actually separated")
# }}}
if sources_are_targets:
# {{{ sep_{smaller,bigger} are duals of each other
assert (trav.target_or_target_parent_boxes == np.arange(
tree.nboxes)).all()
# {{{ list 4 <= list 3
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box + 2]
for jbox in ssn.lists[start:end]:
rstart, rend = trav.sep_bigger_starts[jbox:jbox + 2]
assert ibox in trav.sep_bigger_lists[rstart:rend], (ibox,
jbox)
# }}}
# {{{ list 4 <= list 3
box_to_target_box_index = np.empty(tree.nboxes, tree.box_id_dtype)
box_to_target_box_index.fill(-1)
box_to_target_box_index[trav.target_boxes] = np.arange(
len(trav.target_boxes), dtype=tree.box_id_dtype)
assert (trav.source_boxes == trav.target_boxes).all()
assert (trav.target_or_target_parent_boxes == np.arange(
tree.nboxes, dtype=tree.box_id_dtype)).all()
for ibox in range(tree.nboxes):
start, end = trav.sep_bigger_starts[ibox:ibox + 2]
for jbox in trav.sep_bigger_lists[start:end]:
# In principle, entries of sep_bigger_lists are
# source boxes. In this special case, source and target boxes
# are the same thing (i.e. leaves--see assertion above), so we
# may treat them as targets anyhow.
jtgt_box = box_to_target_box_index[jbox]
assert jtgt_box != -1
good = False
for ssn in trav.sep_smaller_by_level:
rstart, rend = ssn.starts[jtgt_box:jtgt_box + 2]
good = good or ibox in ssn.lists[rstart:rend]
if not good:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
plotter.draw_box(ibox, facecolor='green', alpha=0.5)
plotter.draw_box(jbox, facecolor='red', alpha=0.5)
import matplotlib.pyplot as pt
pt.gca().set_aspect("equal")
pt.show()
# This assertion failing means that ibox's list 4 contains a box
# 'jbox' whose list 3 does not contain ibox.
assert good, (ibox, jbox)
# }}}
logger.info("list 3, 4 are duals")
# }}}
# {{{ sep_smaller satisfies relative level assumption
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box + 2]
for jbox in ssn.lists[start:end]:
assert levels[ibox] < levels[jbox]
logger.info("list 3 satisfies relative level assumption")
# }}}
# {{{ sep_bigger satisfies relative level assumption
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_bigger_starts[itgt_box:itgt_box + 2]
for jbox in trav.sep_bigger_lists[start:end]:
assert levels[tgt_ibox] > levels[jbox]
logger.info("list 4 satisfies relative level assumption")
# }}}
# {{{ level_start_*_box_nrs lists make sense
for name, ref_array in [("level_start_source_box_nrs", trav.source_boxes),
("level_start_source_parent_box_nrs",
trav.source_parent_boxes),
("level_start_target_box_nrs", trav.target_boxes),
("level_start_target_or_target_parent_box_nrs",
trav.target_or_target_parent_boxes)]:
level_starts = getattr(trav, name)
for lev in range(tree.nlevels):
start, stop = level_starts[lev:lev + 2]
box_nrs = ref_array[start:stop]
assert (tree.box_levels[box_nrs] == lev).all(), name
def test_space_invader_query(ctx_getter, dims, dtype, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
dtype = np.dtype(dtype)
nparticles = 10**5
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
nballs = 10**4
ball_centers = make_normal_particle_array(queue, nballs, dims, dtype)
ball_radii = cl.array.empty(queue, nballs, dtype).fill(0.1)
from boxtree.area_query import (
LeavesToBallsLookupBuilder, SpaceInvaderQueryBuilder)
siqb = SpaceInvaderQueryBuilder(ctx)
# 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(ctx)
siq, _ = siqb(queue, tree, ball_centers, ball_radii)
lbl, _ = lblb(queue, tree, ball_centers, ball_radii)
# get data to host for test
tree = tree.get(queue=queue)
siq = siq.get(queue=queue)
lbl = lbl.get(queue=queue)
ball_centers = np.array([x.get() for x in ball_centers])
ball_radii = ball_radii.get()
# 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_CHILDREN:
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)
def test_source_target_tree(ctx_getter, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 2 * 10**5
ntargets = 3 * 10**5
dtype = np.float64
sources = make_normal_particle_array(queue, nsources, dims, dtype,
seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype,
seed=19)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
pt.show()
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, sources, targets=targets,
max_particles_in_box=10, debug=True)
tree = tree.get(queue=queue)
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
assert (sorted_sources
== unsorted_sources[:, tree.user_source_ids]).all()
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets, dtype=np.intp)
assert (sorted_targets
== unsorted_targets[:, user_target_ids]).all()
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
tol = 1e-15
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low
>= tree.bounding_box[0] - 1e-12*tree.root_extent).all(), ibox
assert (extent_high
<= tree.bounding_box[1] + 1e-12*tree.root_extent).all(), ibox
src_start = tree.box_source_starts[ibox]
tgt_start = tree.box_target_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox]
+ np.sum(tree.box_target_counts_cumul[existing_children])
== tree.box_target_counts_cumul[ibox])
for what, particles in [
("sources", sorted_sources[:,
src_start:src_start+tree.box_source_counts_cumul[ibox]]),
("targets", sorted_targets[:,
tgt_start:tgt_start+tree.box_target_counts_cumul[ibox]]),
]:
good = (
(particles < extent_high[:, np.newaxis] + tol)
& (extent_low[:, np.newaxis] - tol <= particles)
).all(axis=0)
all_good_here = good.all()
if do_plot and not all_good_here:
pt.plot(
particles[0, np.where(~good)[0]],
particles[1, np.where(~good)[0]], "ro")
plotter.draw_box(ibox, edgecolor="red")
pt.show()
if not all_good_here:
print("BAD BOX %s %d" % (what, ibox))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_so_far
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
def test_tree_connectivity(ctx_getter, dims, sources_are_targets):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
dtype = np.float64
sources = make_normal_particle_array(queue, 1 * 10**5, dims, dtype)
if sources_are_targets:
targets = None
else:
targets = make_normal_particle_array(queue, 2 * 10**5, dims, dtype)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, max_particles_in_box=30,
targets=targets, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree, debug=True)
tree = tree.get(queue=queue)
trav = trav.get(queue=queue)
levels = tree.box_levels
parents = tree.box_parent_ids.T
children = tree.box_child_ids.T
centers = tree.box_centers.T
# {{{ parent and child relations, levels match up
for ibox in range(1, tree.nboxes):
# /!\ Not testing box 0, has no parents
parent = parents[ibox]
assert levels[parent] + 1 == levels[ibox]
assert ibox in children[parent], ibox
# }}}
if 0:
import matplotlib.pyplot as pt
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.show()
# {{{ neighbor_source_boxes (list 1) consists of source boxes
for itgt_box, ibox in enumerate(trav.target_boxes):
start, end = trav.neighbor_source_boxes_starts[itgt_box:itgt_box+2]
nbl = trav.neighbor_source_boxes_lists[start:end]
if sources_are_targets:
assert ibox in nbl
for jbox in nbl:
assert (0 == children[jbox]).all(), (ibox, jbox, children[jbox])
logger.info("list 1 consists of source boxes")
# }}}
# {{{ separated siblings (list 2) are actually separated
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_siblings_starts[itgt_box:itgt_box+2]
seps = trav.sep_siblings_lists[start:end]
assert (levels[seps] == levels[tgt_ibox]).all()
# three-ish box radii (half of size)
mindist = 2.5 * 0.5 * 2**-int(levels[tgt_ibox]) * tree.root_extent
icenter = centers[tgt_ibox]
for jbox in seps:
dist = la.norm(centers[jbox]-icenter)
assert dist > mindist, (dist, mindist)
logger.info("separated siblings (list 2) are actually separated")
# }}}
if sources_are_targets:
# {{{ sep_{smaller,bigger} are duals of each other
assert (trav.target_or_target_parent_boxes == np.arange(tree.nboxes)).all()
# {{{ list 4 <= list 3
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box+2]
for jbox in ssn.lists[start:end]:
rstart, rend = trav.sep_bigger_starts[jbox:jbox+2]
assert ibox in trav.sep_bigger_lists[rstart:rend], (ibox, jbox)
# }}}
# {{{ list 4 <= list 3
box_to_target_box_index = np.empty(tree.nboxes, tree.box_id_dtype)
box_to_target_box_index.fill(-1)
box_to_target_box_index[trav.target_boxes] = np.arange(
len(trav.target_boxes), dtype=tree.box_id_dtype)
assert (trav.source_boxes == trav.target_boxes).all()
assert (trav.target_or_target_parent_boxes == np.arange(
tree.nboxes, dtype=tree.box_id_dtype)).all()
for ibox in range(tree.nboxes):
start, end = trav.sep_bigger_starts[ibox:ibox+2]
for jbox in trav.sep_bigger_lists[start:end]:
# In principle, entries of sep_bigger_lists are
# source boxes. In this special case, source and target boxes
# are the same thing (i.e. leaves--see assertion above), so we
# may treat them as targets anyhow.
jtgt_box = box_to_target_box_index[jbox]
assert jtgt_box != -1
good = False
for ssn in trav.sep_smaller_by_level:
rstart, rend = ssn.starts[jtgt_box:jtgt_box+2]
good = good or ibox in ssn.lists[rstart:rend]
if not good:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
plotter.draw_box(ibox, facecolor='green', alpha=0.5)
plotter.draw_box(jbox, facecolor='red', alpha=0.5)
import matplotlib.pyplot as pt
pt.gca().set_aspect("equal")
pt.show()
# This assertion failing means that ibox's list 4 contains a box
# 'jbox' whose list 3 does not contain ibox.
assert good, (ibox, jbox)
# }}}
logger.info("list 3, 4 are duals")
# }}}
# {{{ sep_smaller satisfies relative level assumption
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box+2]
for jbox in ssn.lists[start:end]:
assert levels[ibox] < levels[jbox]
logger.info("list 3 satisfies relative level assumption")
# }}}
# {{{ sep_bigger satisfies relative level assumption
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_bigger_starts[itgt_box:itgt_box+2]
for jbox in trav.sep_bigger_lists[start:end]:
assert levels[tgt_ibox] > levels[jbox]
logger.info("list 4 satisfies relative level assumption")
# }}}
# {{{ level_start_*_box_nrs lists make sense
for name, ref_array in [
("level_start_source_box_nrs", trav.source_boxes),
("level_start_source_parent_box_nrs", trav.source_parent_boxes),
("level_start_target_box_nrs", trav.target_boxes),
("level_start_target_or_target_parent_box_nrs",
trav.target_or_target_parent_boxes)
]:
level_starts = getattr(trav, name)
for lev in range(tree.nlevels):
start, stop = level_starts[lev:lev+2]
box_nrs = ref_array[start:stop]
assert (tree.box_levels[box_nrs] == lev).all(), name
def test_extent_tree(ctx_getter, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 100000
ntargets = 200000
dtype = np.float64
npoint_sources_per_source = 16
sources = make_normal_particle_array(queue, nsources, dims, dtype,
seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype,
seed=19)
from pyopencl.clrandom import RanluxGenerator
rng = RanluxGenerator(queue, seed=13)
source_radii = 2**rng.uniform(queue, nsources, dtype=dtype,
a=-10, b=0)
target_radii = 2**rng.uniform(queue, ntargets, dtype=dtype,
a=-10, b=0)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
dev_tree, _ = tb(queue, sources, targets=targets,
source_radii=source_radii, target_radii=target_radii,
max_particles_in_box=10, debug=True)
logger.info("transfer tree, check orderings")
tree = dev_tree.get(queue=queue)
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
sorted_source_radii = tree.source_radii
sorted_target_radii = tree.target_radii
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
unsorted_source_radii = source_radii.get()
unsorted_target_radii = target_radii.get()
assert (sorted_sources
== unsorted_sources[:, tree.user_source_ids]).all()
assert (sorted_source_radii
== unsorted_source_radii[tree.user_source_ids]).all()
# {{{ test box structure, stick-out criterion
logger.info("test box structure, stick-out criterion")
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets, dtype=np.intp)
if ntargets:
assert (sorted_targets
== unsorted_targets[:, user_target_ids]).all()
assert (sorted_target_radii
== unsorted_target_radii[user_target_ids]).all()
all_good_so_far = True
# {{{ check sources, targets
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
box_radius = np.max(extent_high-extent_low) * 0.5
stick_out_dist = tree.stick_out_factor * box_radius
assert (extent_low >=
tree.bounding_box[0] - 1e-12*tree.root_extent).all(), ibox
assert (extent_high <=
tree.bounding_box[1] + 1e-12*tree.root_extent).all(), ibox
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox]
+ np.sum(tree.box_target_counts_cumul[existing_children])
== tree.box_target_counts_cumul[ibox])
for what, starts, counts, points, radii in [
("source", tree.box_source_starts, tree.box_source_counts_cumul,
sorted_sources, sorted_source_radii),
("target", tree.box_target_starts, tree.box_target_counts_cumul,
sorted_targets, sorted_target_radii),
]:
bstart = starts[ibox]
bslice = slice(bstart, bstart+counts[ibox])
check_particles = points[:, bslice]
check_radii = radii[bslice]
good = (
(check_particles + check_radii
< extent_high[:, np.newaxis] + stick_out_dist)
&
(extent_low[:, np.newaxis] - stick_out_dist
<= check_particles - check_radii)
).all(axis=0)
all_good_here = good.all()
if not all_good_here:
print("BAD BOX %s %d level %d" % (what, ibox, tree.box_levels[ibox]))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_here
# }}}
assert all_good_so_far
# }}}
# {{{ create, link point sources
logger.info("creating point sources")
np.random.seed(20)
from pytools.obj_array import make_obj_array
point_sources = make_obj_array([
cl.array.to_device(queue,
unsorted_sources[i][:, np.newaxis]
+ unsorted_source_radii[:, np.newaxis]
* np.random.uniform(
-1, 1, size=(nsources, npoint_sources_per_source))
)
for i in range(dims)])
point_source_starts = cl.array.arange(queue,
0, (nsources+1)*npoint_sources_per_source, npoint_sources_per_source,
dtype=tree.particle_id_dtype)
from boxtree.tree import link_point_sources
dev_tree = link_point_sources(queue, dev_tree,
point_source_starts, point_sources,
debug=True)
