diff --git a/boxtree/area_query.py b/boxtree/area_query.py new file mode 100644 index 0000000000000000000000000000000000000000..63d6fd2c36f5afa7406db132522efb27f82cae93 --- /dev/null +++ b/boxtree/area_query.py @@ -0,0 +1,561 @@ +# -*- coding: utf-8 -*- +from __future__ import division + +__copyright__ = """Copyright (C) 2016 Matt Wala""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +import numpy as np +import pyopencl as cl +import pyopencl.array # noqa +from mako.template import Template +from boxtree.tools import AXIS_NAMES, DeviceDataRecord +from pytools import memoize_method + +import logging +logger = logging.getLogger(__name__) + + +# {{{ output + +class PeerListLookup(DeviceDataRecord): + """ + .. attribute:: tree + + The :class:`boxtree.Tree` instance used to build this lookup. + + .. attribute:: peer_list_starts + + Indices into :attr:`peer_lists`. + ``peer_lists[peer_list_starts[box_id]:peer_list_starts[box_id]+1]`` + contains the list of peer boxes of box `box_id`. + + .. attribute:: peer_lists + """ + + +class AreaQuery(DeviceDataRecord): + """ + .. attribute:: tree + + The :class:`boxtree.Tree` instance used to build this lookup. + + .. attribute:: leaves_near_ball_starts + + Indices into :attr:`query_lists`. + ``query_lists[leaves_near_ball_starts[ball_nr]: + leaves_near_ball_starts[ball_nr]+1]`` + results in a list of leaf boxes that intersect `ball_nr`. + + .. attribute:: leaves_near_ball_lists + """ + +# }}} + +# {{{ kernel templates + +AREA_QUERY_TEMPLATE = r"""//CL// +typedef ${dtype_to_ctype(ball_id_dtype)} ball_id_t; +typedef ${dtype_to_ctype(peer_list_idx_dtype)} peer_list_idx_t; + +<%def name="add_box_to_list_if_overlaps_ball(box_id)"> + { + bool is_overlapping; + + ${check_ball_overlap("is_overlapping", box_id, "ball_radius", "ball_center")} + + if (is_overlapping) + { + APPEND_leaves(${box_id}); + } + } + + +void generate(LIST_ARG_DECL USER_ARG_DECL ball_id_t ball_nr) +{ + coord_vec_t ball_center; + %for i in range(dimensions): + ball_center.${AXIS_NAMES[i]} = ball_${AXIS_NAMES[i]}[ball_nr]; + %endfor + + coord_t ball_radius = ball_radii[ball_nr]; + + /////////////////////////////////// + // Step 1: Find the guiding box. // + /////////////////////////////////// + + ${walk_init(0)} + box_id_t guiding_box; + + if (LEVEL_TO_RAD(0) < ball_radius / 2 || !(box_flags[0] & BOX_HAS_CHILDREN)) + { + guiding_box = 0; + continue_walk = false; + } + + while (continue_walk) + { + // Get the next child. + box_id_t child_box_id = box_child_ids[ + walk_morton_nr * aligned_nboxes + walk_box_id]; + + bool last_child = walk_morton_nr == ${2**dimensions - 1}; + + if (child_box_id) + { + bool contains_ball_center; + int child_level = walk_level + 1; + coord_t child_rad = LEVEL_TO_RAD(child_level); + + { + // Check if the child contains the ball's center. + ${load_center("child_center", "child_box_id")} + + coord_t max_dist = 0; + %for i in range(dimensions): + max_dist = fmax(max_dist, + fabs(ball_center.s${i} - child_center.s${i})); + %endfor + + contains_ball_center = max_dist <= child_rad; + } + + if (contains_ball_center) + { + if ((child_rad / 2 < ball_radius && ball_radius <= child_rad) || + !(box_flags[child_box_id] & BOX_HAS_CHILDREN)) + { + guiding_box = child_box_id; + break; + } + + // We want to descend into this box. Put the current state + // on the stack. + ${walk_push("child_box_id")} + continue; + } + } + + if (last_child) + { + // This box has no children that contain the center, so it must + // be the guiding box. + guiding_box = walk_box_id; + break; + } + + ${walk_advance()} + } + + ////////////////////////////////////////////////////// + // Step 2 - Walk the peer boxes to find the leaves. // + ////////////////////////////////////////////////////// + + + for (peer_list_idx_t pb_i = peer_list_starts[guiding_box], + pb_e = peer_list_starts[guiding_box+1]; pb_i < pb_e; ++pb_i) + { + box_id_t peer_box = peer_lists[pb_i]; + + if (!(box_flags[peer_box] & BOX_HAS_CHILDREN)) + { + ${add_box_to_list_if_overlaps_ball("peer_box")} + } + else + { + ${walk_reset("peer_box")} + + while (continue_walk) + { + box_id_t child_box_id = box_child_ids[ + walk_morton_nr * aligned_nboxes + walk_box_id]; + + if (child_box_id) + { + if (!(box_flags[child_box_id] & BOX_HAS_CHILDREN)) + { + ${add_box_to_list_if_overlaps_ball("child_box_id")} + } + else + { + // We want to descend into this box. Put the current state + // on the stack. + ${walk_push("child_box_id")} + continue; + } + } + + ${walk_advance()} + } + } + } +} +""" + +PEER_LIST_FINDER_TEMPLATE = r"""//CL// + +void generate(LIST_ARG_DECL USER_ARG_DECL box_id_t box_id) +{ + ${load_center("center", "box_id")} + + if (box_id == 0) + { + // Peer of root = self + APPEND_peers(box_id); + return; + } + + int level = box_levels[box_id]; + + // To find this box's peers, start at the top of the tree, descend + // into adjacent (or overlapping) parents. + ${walk_init(0)} + + while (continue_walk) + { + box_id_t child_box_id = box_child_ids[ + walk_morton_nr * aligned_nboxes + walk_box_id]; + + if (child_box_id) + { + ${load_center("child_center", "child_box_id")} + + // child_box_id lives on walk_level+1. + bool a_or_o = is_adjacent_or_overlapping(root_extent, + center, level, child_center, walk_level+1, false); + + if (a_or_o) + { + // child_box_id lives on walk_level+1. + if (walk_level+1 == level) + { + APPEND_peers(child_box_id); + } + else if (!(box_flags[child_box_id] & BOX_HAS_CHILDREN)) + { + APPEND_peers(child_box_id); + } + else + { + // Check if any children are adjacent or overlapping. + // If not, this box must be a peer. + bool must_be_peer = true; + + for (int morton_nr = 0; + must_be_peer && morton_nr < ${2**dimensions}; + ++morton_nr) + { + box_id_t next_child_id = box_child_ids[ + morton_nr * aligned_nboxes + child_box_id]; + if (next_child_id) + { + ${load_center("next_child_center", "next_child_id")} + must_be_peer &= !is_adjacent_or_overlapping(root_extent, + center, level, next_child_center, walk_level+2, + false); + } + } + + if (must_be_peer) + { + APPEND_peers(child_box_id); + } + else + { + // We want to descend into this box. Put the current state + // on the stack. + ${walk_push("child_box_id")} + continue; + } + } + } + } + + ${walk_advance()} + } +} + +""" + +# }}} + + +class AreaQueryBuilder(object): + """Given a set of :math:`l^\infty` "balls", this class helps build a + look-up table from ball to leaf boxes that intersect with the ball. + """ + def __init__(self, context): + self.context = context + + # {{{ Kernel generation + + @memoize_method + def get_area_query_kernel(self, dimensions, coord_dtype, box_id_dtype, + ball_id_dtype, peer_list_idx_dtype, max_levels): + from pyopencl.tools import dtype_to_ctype + from boxtree import box_flags_enum + + logger.info("start building area query kernel") + + from boxtree.traversal import TRAVERSAL_PREAMBLE_TEMPLATE + + template = Template( + TRAVERSAL_PREAMBLE_TEMPLATE + + AREA_QUERY_TEMPLATE, + strict_undefined=True) + + render_vars = dict( + dimensions=dimensions, + dtype_to_ctype=dtype_to_ctype, + box_id_dtype=box_id_dtype, + particle_id_dtype=None, + coord_dtype=coord_dtype, + vec_types=cl.array.vec.types, + max_levels=max_levels, + AXIS_NAMES=AXIS_NAMES, + box_flags_enum=box_flags_enum, + peer_list_idx_dtype=peer_list_idx_dtype, + ball_id_dtype=ball_id_dtype, + debug=False, + # Not used + stick_out_factor=0) + + from pyopencl.tools import VectorArg, ScalarArg + arg_decls = [ + VectorArg(coord_dtype, "box_centers"), + ScalarArg(coord_dtype, "root_extent"), + VectorArg(np.uint8, "box_levels"), + ScalarArg(box_id_dtype, "aligned_nboxes"), + VectorArg(box_id_dtype, "box_child_ids"), + VectorArg(box_flags_enum.dtype, "box_flags"), + VectorArg(peer_list_idx_dtype, "peer_list_starts"), + VectorArg(box_id_dtype, "peer_lists"), + VectorArg(coord_dtype, "ball_radii"), + ] + [ + VectorArg(coord_dtype, "ball_"+ax) + for ax in AXIS_NAMES[:dimensions]] + + from pyopencl.algorithm import ListOfListsBuilder + area_query_kernel = ListOfListsBuilder( + self.context, + [("leaves", box_id_dtype)], + str(template.render(**render_vars)), + arg_decls=arg_decls, + name_prefix="area_query", + count_sharing={}, + complex_kernel=True) + + logger.info("done building area query kernel") + return area_query_kernel + + # }}} + + @memoize_method + def get_peer_list_finder(self): + return PeerListFinder(self.context) + + def __call__(self, queue, tree, ball_centers, ball_radii, peer_lists=None, + wait_for=None): + """ + :arg queue: a :class:`pyopencl.CommandQueue` + :arg tree: a :class:`boxtree.Tree`. + :arg ball_centers: an object array of coordinate + :class:`pyopencl.array.Array` instances. + Their *dtype* must match *tree*'s + :attr:`boxtree.Tree.coord_dtype`. + :arg ball_radii: a + :class:`pyopencl.array.Array` + of positive numbers. + Its *dtype* must match *tree*'s + :attr:`boxtree.Tree.coord_dtype`. + :arg peer_lists: may either be *None* or an instance of + :class:`PeerListLookup` associated with `tree`. + :arg wait_for: may either be *None* or a list of :class:`pyopencl.Event` + instances for whose completion this command waits before starting + exeuction. + :returns: a tuple *(aq, event)*, where *lbl* is an instance of + :class:`AreaQuery`, and *event* is a :class:`pyopencl.Event` + for dependency management. + """ + + from pytools import single_valued + if single_valued(bc.dtype for bc in ball_centers) != tree.coord_dtype: + raise TypeError("ball_centers dtype must match tree.coord_dtype") + if ball_radii.dtype != tree.coord_dtype: + raise TypeError("ball_radii dtype must match tree.coord_dtype") + + ball_id_dtype = tree.particle_id_dtype # ? + + from pytools import div_ceil + # Avoid generating too many kernels. + max_levels = div_ceil(tree.nlevels, 10) * 10 + + if peer_lists is None: + peer_list_finder = self.get_peer_list_finder() + peer_lists, evt = peer_list_finder(queue, tree, wait_for=wait_for) + wait_for = [evt] + + if len(peer_lists.peer_list_starts) != tree.nboxes + 1: + raise ValueError("size of peer lists must match with number of boxes") + + area_query_kernel = self.get_area_query_kernel(tree.dimensions, + tree.coord_dtype, tree.box_id_dtype, ball_id_dtype, + peer_lists.peer_list_starts.dtype, max_levels) + + logger.info("area query: run area query") + + result, evt = area_query_kernel( + queue, len(ball_radii), + tree.box_centers.data, tree.root_extent, + tree.box_levels.data, tree.aligned_nboxes, + tree.box_child_ids.data, tree.box_flags.data, + peer_lists.peer_list_starts.data, + peer_lists.peer_lists.data, ball_radii.data, + *tuple(bc.data for bc in ball_centers), + wait_for=wait_for) + + logger.info("area query: done") + + return AreaQuery( + tree=tree, + leaves_near_ball_starts=result["leaves"].starts, + leaves_near_ball_lists=result["leaves"].lists).with_queue(None), evt + + +class PeerListFinder(object): + """This class builds a look-up table from box numbers to peer boxes. The + full definition [1]_ of a peer box is as follows: + + Given a box :math:`b_j` in a quad-tree, :math:`b_k` is a peer box of + :math:`b_j` if it is + + 1. adjacent to :math:`b_j`, + + 2. of at least the same size as :math:`b_j` (i.e. at the same or a + higher level than), and + + 3. no child of :math:`b_k` satisfies the above two criteria. + + .. [1] Rachh, Manas, Andreas Klöckner, and Michael O'Neil. "Fast + algorithms for Quadrature by Expansion I: Globally valid expansions." + """ + + def __init__(self, context): + self.context = context + + # {{{ Kernel generation + + @memoize_method + def get_peer_list_finder_kernel(self, dimensions, coord_dtype, + box_id_dtype, max_levels): + from pyopencl.tools import dtype_to_ctype + from boxtree import box_flags_enum + + logger.info("start building peer list finder kernel") + + from boxtree.traversal import ( + TRAVERSAL_PREAMBLE_TEMPLATE, HELPER_FUNCTION_TEMPLATE) + + template = Template( + TRAVERSAL_PREAMBLE_TEMPLATE + + HELPER_FUNCTION_TEMPLATE + + PEER_LIST_FINDER_TEMPLATE, + strict_undefined=True) + + render_vars = dict( + dimensions=dimensions, + dtype_to_ctype=dtype_to_ctype, + box_id_dtype=box_id_dtype, + particle_id_dtype=None, + coord_dtype=coord_dtype, + vec_types=cl.array.vec.types, + max_levels=max_levels, + AXIS_NAMES=AXIS_NAMES, + box_flags_enum=box_flags_enum, + debug=False, + # Not used + stick_out_factor=0, + # For calls to the helper is_adjacent_or_overlapping() + targets_have_extent=False, + sources_have_extent=False) + + from pyopencl.tools import VectorArg, ScalarArg + arg_decls = [ + VectorArg(coord_dtype, "box_centers"), + ScalarArg(coord_dtype, "root_extent"), + VectorArg(np.uint8, "box_levels"), + ScalarArg(box_id_dtype, "aligned_nboxes"), + VectorArg(box_id_dtype, "box_child_ids"), + VectorArg(box_flags_enum.dtype, "box_flags"), + ] + + from pyopencl.algorithm import ListOfListsBuilder + peer_list_finder_kernel = ListOfListsBuilder( + self.context, + [("peers", box_id_dtype)], + str(template.render(**render_vars)), + arg_decls=arg_decls, + name_prefix="find_peer_lists", + count_sharing={}, + complex_kernel=True) + + logger.info("done building peer list finder kernel") + return peer_list_finder_kernel + + # }}} + + def __call__(self, queue, tree, wait_for=None): + """ + :arg queue: a :class:`pyopencl.CommandQueue` + :arg tree: a :class:`boxtree.Tree`. + :arg wait_for: may either be *None* or a list of :class:`pyopencl.Event` + instances for whose completion this command waits before starting + execution. + :returns: a tuple *(pl, event)*, where *pl* is an instance of + :class:`PeerListLookup`, and *event* is a :class:`pyopencl.Event` + for dependency management. + """ + from pytools import div_ceil + # Avoid generating too many kernels. + max_levels = div_ceil(tree.nlevels, 10) * 10 + + peer_list_finder_kernel = self.get_peer_list_finder_kernel( + tree.dimensions, tree.coord_dtype, tree.box_id_dtype, max_levels) + + logger.info("peer list finder: find peer lists") + + result, evt = peer_list_finder_kernel( + queue, tree.nboxes, + tree.box_centers.data, tree.root_extent, + tree.box_levels.data, tree.aligned_nboxes, + tree.box_child_ids.data, tree.box_flags.data, + wait_for=wait_for) + + logger.info("peer list finder: done") + + return PeerListLookup( + tree=tree, + peer_list_starts=result["peers"].starts, + peer_lists=result["peers"].lists).with_queue(None), evt + + +# vim: filetype=pyopencl:fdm=marker diff --git a/boxtree/geo_lookup.py b/boxtree/geo_lookup.py index ffebd1f147cc8bad830990a8f8aecc65fa544371..d328628d73cebc507503a8563707edf4f1ca07fc 100644 --- a/boxtree/geo_lookup.py +++ b/boxtree/geo_lookup.py @@ -86,20 +86,8 @@ void generate(LIST_ARG_DECL USER_ARG_DECL ball_id_t ball_nr) { bool is_overlapping; - { - ${load_center("child_center", "child_box_id")} - int child_level = box_levels[child_box_id]; - - coord_t size_sum = LEVEL_TO_RAD(child_level) + ball_radius; - - coord_t max_dist = 0; - %for i in range(dimensions): - max_dist = fmax(max_dist, - fabs(ball_center.s${i} - child_center.s${i})); - %endfor - - is_overlapping = max_dist <= size_sum; - } + ${check_ball_overlap( + "is_overlapping", "child_box_id", "ball_radius", "ball_center")} if (is_overlapping) { @@ -229,6 +217,7 @@ class LeavesToBallsLookupBuilder(object): ball_id_dtype = tree.particle_id_dtype # ? from pytools import div_ceil + # Avoid generating too many kernels. max_levels = div_ceil(tree.nlevels, 10) * 10 b2l_knl = self.get_balls_to_leaves_kernel( diff --git a/boxtree/traversal.py b/boxtree/traversal.py index 52584aa524ca23dcb464dbcf490bac6aa85088f4..5d872300f4cee4dda93f62edc3e0ea1ebc7f6dda 100644 --- a/boxtree/traversal.py +++ b/boxtree/traversal.py @@ -77,6 +77,13 @@ typedef ${dtype_to_ctype(vec_types[coord_dtype, dimensions])} coord_vec_t; bool continue_walk = true; +<%def name="walk_reset(start_box_id)"> + walk_level = 0; + walk_box_id = ${start_box_id}; + walk_morton_nr = 0; + continue_walk = true; + + <%def name="walk_advance()"> while (true) { @@ -120,6 +127,23 @@ typedef ${dtype_to_ctype(vec_types[coord_dtype, dimensions])} coord_vec_t; walk_morton_nr = 0; +<%def name="check_ball_overlap(is_overlapping, box_id, ball_radius, ball_center)"> + { + ${load_center("box_center", box_id)} + int box_level = box_levels[${box_id}]; + + coord_t size_sum = LEVEL_TO_RAD(box_level) + ${ball_radius}; + + coord_t max_dist = 0; + %for i in range(dimensions): + max_dist = fmax(max_dist, + fabs(${ball_center}.s${i} - box_center.s${i})); + %endfor + + ${is_overlapping} = max_dist <= size_sum; + } + + """ # }}} diff --git a/doc/lookup.rst b/doc/lookup.rst index af74ed1bb7d1dbef3057eeb2a55c52d722b1ca65..3fd1d9dc1b4f832ab1c6b4b938a3c5c03cb68c96 100644 --- a/doc/lookup.rst +++ b/doc/lookup.rst @@ -11,4 +11,24 @@ Tree-based geometric lookup .. automethod:: get +.. module:: boxtree.area_query + +.. autoclass:: AreaQueryBuilder + + .. automethod:: __call__ + +.. autoclass:: AreaQuery + + .. automethod:: get + +Area queries are implemented using peer lists. + +.. autoclass:: PeerListFinder + + .. automethod:: __call__ + +.. autoclass:: PeerListLookup + + .. automethod:: get + .. vim: sw=4 diff --git a/test/test_tree.py b/test/test_tree.py index b3d1841d5e5b7216ecaa39993e712e6394c4068e..a02d59c6758d62080072f82b4add0455ce62ea38 100644 --- a/test/test_tree.py +++ b/test/test_tree.py @@ -610,6 +610,75 @@ def test_geometry_query(ctx_getter, dims, do_plot=False): # }}} +# {{{ area query test + +@pytest.mark.opencl +@pytest.mark.parametrize("dims", [2, 3]) +def test_area_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 AreaQueryBuilder + aqb = AreaQueryBuilder(ctx) + + area_query, _ = aqb(queue, tree, ball_centers, ball_radii) + + # Get data to host for test. + tree = tree.get(queue=queue) + area_query = area_query.get(queue=queue) + ball_centers = np.array([x.get() for x in ball_centers]).T + ball_radii = ball_radii.get() + + from boxtree import box_flags_enum + + leaf_boxes = np.array([ibox for ibox in range(tree.nboxes) + if tree.box_flags[ibox] & ~box_flags_enum.HAS_CHILDREN]) + + leaf_box_radii = np.empty(len(leaf_boxes)) + leaf_box_centers = np.empty((len(leaf_boxes), dims)) + + for idx, leaf_box in enumerate(leaf_boxes): + box_center = tree.box_centers[:, leaf_box] + ext_l, ext_h = tree.get_box_extent(leaf_box) + leaf_box_radii[idx] = np.max(ext_h - ext_l) * 0.5 + leaf_box_centers[idx] = box_center + + for ball_nr, (ball_center, ball_radius) \ + in enumerate(zip(ball_centers, ball_radii)): + linf_box_dists = np.max(np.abs(ball_center - leaf_box_centers), axis=-1) + near_leaves_indices, \ + = np.where(linf_box_dists < ball_radius + leaf_box_radii) + near_leaves = leaf_boxes[near_leaves_indices] + + start, end = area_query.leaves_near_ball_starts[ball_nr:ball_nr+2] + found = area_query.leaves_near_ball_lists[start:end] + actual = near_leaves + assert sorted(found) == sorted(actual) + +# }}} + + # You can test individual routines by typing # $ python test_tree.py 'test_routine(cl.create_some_context)'