def space_invader_query(self): from boxtree.area_query import SpaceInvaderQueryBuilder return SpaceInvaderQueryBuilder(self.cl_context)
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)