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 inspect_geo_data(insn, bound_expr, geo_data):
nonlocal sizes, nsources, ncenters
tree = geo_data.tree().with_queue(queue)
from boxtree.area_query import PeerListFinder
plf = PeerListFinder(queue.context)
pl, evt = plf(queue, tree)
# Perform an area query around each QBX center, counting the
# neighborhood sizes.
knl = NeighborhoodCounter.generate(
queue.context,
tree.dimensions,
tree.coord_dtype,
tree.box_id_dtype,
tree.box_id_dtype,
tree.nlevels,
extra_type_aliases=(('particle_id_t', tree.particle_id_dtype),))
centers = geo_data.centers()
search_radii = radius * geo_data.expansion_radii().with_queue(queue)
ncenters = len(search_radii)
nsources = tree.nsources
sizes = cl.array.zeros(queue, ncenters, np.int32)
assert nsources == lpot_source.quad_stage2_density_discr.nnodes
coords = []
coords.extend(tree.sources)
coords.extend(centers)
evt = knl(
*NeighborhoodCounter.unwrap_args(
tree,
pl,
tree.box_source_starts,
tree.box_source_counts_cumul,
search_radii,
sizes,
*coords),
range=slice(ncenters),
queue=queue,
wait_for=[evt])
cl.wait_for_events([evt])
return False # no need to do the actual FMM
def peer_list_finder(self):
from boxtree.area_query import PeerListFinder
return PeerListFinder(self.array_context.context)