From 11d1ffbcce0affffad9abaf752d4b3d6deb8f1a0 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 12 May 2017 15:29:46 -0500 Subject: [PATCH] Fix leaves to balls lookup size (closes #10 on gitlab). --- boxtree/area_query.py | 18 +++++++++++++----- test/test_tree.py | 2 ++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/boxtree/area_query.py b/boxtree/area_query.py index 88cc4b6..a5ebdca 100644 --- a/boxtree/area_query.py +++ b/boxtree/area_query.py @@ -835,16 +835,24 @@ class LeavesToBallsLookupBuilder(object): logger.info("leaves-to-balls lookup: expand starts") - nkeys = len(area_query.leaves_near_ball_lists) + nkeys = tree.nboxes nballs_p_1 = len(area_query.leaves_near_ball_starts) assert nballs_p_1 == len(ball_radii) + 1 + # We invert the area query in two steps: + # + # 1. Turn the area query result into (ball number, box number) pairs. + # This is done in the "starts expander kernel." + # + # 2. Key-value sort the (ball number, box number) pairs by box number. + starts_expander_knl = self.get_starts_expander_kernel(tree.box_id_dtype) - expanded_starts = cl.array.empty(queue, nkeys, tree.box_id_dtype) + expanded_starts = cl.array.empty( + queue, len(area_query.leaves_near_ball_lists), tree.box_id_dtype) evt = starts_expander_knl( - expanded_starts, - area_query.leaves_near_ball_starts.with_queue(queue), - nballs_p_1) + expanded_starts, + area_query.leaves_near_ball_starts.with_queue(queue), + nballs_p_1) wait_for = [evt] logger.info("leaves-to-balls lookup: key-value sort") diff --git a/test/test_tree.py b/test/test_tree.py index 9fc0833..0857787 100644 --- a/test/test_tree.py +++ b/test/test_tree.py @@ -636,6 +636,8 @@ def test_leaves_to_balls_query(ctx_getter, dims, do_plot=False): 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): -- GitLab