diff --git a/boxtree/tree_build.py b/boxtree/tree_build.py index cf4e80b077e53e0eb4a3797b274556a255bc3334..e20290142c7bbbc57b49549ba728c85f6d41e21e 100644 --- a/boxtree/tree_build.py +++ b/boxtree/tree_build.py @@ -123,6 +123,8 @@ class TreeBuilder(object): execution. :arg extent_norm: ``"l2"`` or ``"linf"``. Indicates the norm with respect to which particle stick-out is measured. See :attr:`Tree.extent_norm`. + :arg bbox: Bounding box of the same type as returned by + *boxtree.bounding_box.make_bounding_box_dtype*. :arg bbox: Bounding box of either type: 1. A dim-by-2 array, with each row to be [min, max] coordinates in its corresponding axis direction. @@ -132,8 +134,8 @@ class TreeBuilder(object): building. Otherwise, the bounding box is determined from particles in such a way that it is square and is slightly larger at the top (so that scaled coordinates are always < 1). - When supplied, the bounding box must be square and have all the - particles in its closure. + When given, this bounding box is used for tree + building. Otherwise, the bounding box is determined from particles. :arg kwargs: Used internally for debugging. :returns: a tuple ``(tree, event)``, where *tree* is an instance of diff --git a/boxtree/tree_interactive_build.py b/boxtree/tree_interactive_build.py new file mode 100644 index 0000000000000000000000000000000000000000..63e0bc6c8c9ba8e0d4dc441dc95352d360e4a938 --- /dev/null +++ b/boxtree/tree_interactive_build.py @@ -0,0 +1,660 @@ +from __future__ import division, absolute_import + +__copyright__ = "Copyright (C) 2018 Xiaoyu Wei" + +__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 +from cgen import Enum +from itertools import product + +import loopy as lp +from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_2 # noqa +import pyopencl as cl +import pyopencl.array # noqa +from pytools.obj_array import make_obj_array +from boxtree.tools import DeviceDataRecord + +# {{{ adaptivity_flags + +class adaptivity_flags_enum(Enum): # noqa + """Constants for box adaptivity flags bit field.""" + + c_name = "adaptivity_flags_t" + dtype = np.dtype(np.uint8) + c_value_prefix = "ADAPTIVITY_FLAG_" + + REFINE = 1 << 0 + COARSEN = 1 << 1 + +# }}} + +# {{{ Data structure for a tree of boxes + + +class BoxTree(DeviceDataRecord): + r"""A quad/oct-tree tree consisting of a hierarchy of boxes. + The intended use for this data structure is to generate particles + as quadrature points in each leaf box. Those particles can then + be used to build a Tree object and then generate the Traversal for + FMM. + + This class is designed such that it is easy to adaptively modify + the tree structure (refine/coarsen). Unlike flattened Tree class, + it handles itself and does not rely on external tree builders. + + .. ------------------------------------------------------------------------ + .. rubric:: Data types + .. ------------------------------------------------------------------------ + + .. attribute:: box_id_dtype + .. attribute:: coord_dtype + .. attribute:: box_level_dtype + + .. ------------------------------------------------------------------------ + .. rubric:: Counts and sizes + .. ------------------------------------------------------------------------ + + .. attribute:: root_extent + + the root box size, a scalar + + .. attribute:: nlevels + + .. attribute:: nboxes + + Can be larger than the actual number of boxes, since the box ids of + purged boxes during coarsening are not reused. + + .. attribute:: nboxes_level + + ``size_t [nlevels]`` + + Can be larger than the actual number of boxes at each level due to + the same reason for nboxes. + + .. attribute:: n_active_boxes + + .. ------------------------------------------------------------------------ + .. rubric:: Box properties + .. ------------------------------------------------------------------------ + + .. attribute:: box_centers + + ``coord_t [dimensions, nboxes]`` + + (C order, 'structure of arrays') + + .. attribute:: box_levels + + :attr:`box_level_dtype` ``box_level_t [nboxes]`` + + .. attribute:: box_is_active + + :attr:`bool` ``bool [nboxes]`` + + FIXME: pyopencl cannot map 'bool'. Resort to use an int for now. + + .. ------------------------------------------------------------------------ + .. rubric:: Structural properties + .. ------------------------------------------------------------------------ + + .. attribute:: level_boxes + + ``numpy.ndarray [nlevels]`` + ``box_id_t [nlevels][nboxes_level[level]]`` + + A :class:`numpy.ndarray` of box ids at each level. It acts as an + inverse lookup table of box_levels. The outer layer is an object + array to account for variable lengths at each level. + + .. attribute:: box_parent_ids + + ``box_id_t [nboxes]`` + + Box 0 (the root) has 0 as its parent. + + .. attribute:: box_child_ids + + ``box_id_t [2**dimensions, nboxes]`` + + (C order, 'structure of arrays') + + "0" is used as a 'no child' marker, as the root box can never + occur as any box's child. Boxes with no child are called active + boxes. + + .. attribute:: active_boxes + + ``box_id_t [n_active_boxes]`` + + .. ------------------------------------------------------------------------ + .. rubric:: Adaptivity properties + .. ------------------------------------------------------------------------ + + .. attribute:: adaptivity_flags + + :attr:`adaptivity_flags_enum.dtype` ``[n_active_boxes]`` + + A bitwise combination of :class:`adaptivity_flags_enum` constants. + + .. attribute:: n_active_neighbors + + ``size_t [n_active_boxes]`` + + Number of adjacent active boxes for each active box. + + .. attribute:: neighboring_active_boxes + + ``box_id_t [n_active_boxes, n_active_neighbors[n_active_boxes]]`` + + List of adjacent active boxes for each active box. The root box + cannot be in this list of any box, and 0 is used as a placeholder + in the case that the actual number of neighbors is less than + the value held in n_active_neighbors[]. + """ + def get_copy_kwargs(self, **kwargs): + # cl arrays + for f in self.__class__.fields: + if f not in kwargs: + try: + kwargs[f] = getattr(self, f) + except AttributeError: + pass + + # others + kwargs.update({ + "size_t": self.size_t, + "box_id_dtype": self.box_id_dtype, + "box_level_dtype": self.box_level_dtype, + "coord_dtype": self.coord_dtype, + "root_extent": self.root_extent, + }) + + return kwargs + + def generate_uniform_boxtree(self, queue, + root_vertex=np.zeros(2), + root_extent=1, + nlevels=1, + box_id_dtype=np.int32, + box_level_dtype=np.int32, + coord_dtype=np.float64): + """A plain boxtree with uniform levels (a complete tree). + The root box is given its vertex with the smallest coordinates + and its extent. + """ + self.size_t = np.int32 + self.box_id_dtype = box_id_dtype + self.box_level_dtype = box_level_dtype + self.coord_dtype = coord_dtype + + dim = len(root_vertex) + self.root_extent = root_extent + + self.nboxes_level = cl.array.to_device(queue, + np.array( + [1 << (dim * l) for l in range(nlevels)], + dtype=self.size_t)) + self.register_fields({"nboxes_level": self.nboxes_level}) + + nboxes = self.size_t(cl.array.sum(self.nboxes_level).get()) + + self.box_levels = cl.array.zeros(queue, nboxes, box_level_dtype) + level_start = 0 + for l in range(nlevels): + offset = self.size_t(self.nboxes_level[l].get()) + self.box_levels[level_start: level_start + offset] = l + level_start += offset + self.register_fields({"box_levels": self.box_levels}) + + self.box_centers = cl.array.zeros(queue, (dim, nboxes), coord_dtype) + ibox = 0 + for l in range(nlevels): + dx = self.root_extent / (1 << l) + for cid in product(range(1 << l), repeat=dim): + for d in range(dim): + self.box_centers[d, ibox] = cid[d] * dx + ( + dx / 2 + root_vertex[d]) + ibox += 1 + self.register_fields({"box_centers": self.box_centers}) + + n_active_boxes = self.size_t(self.nboxes_level[nlevels - 1].get()) + self.active_boxes = cl.array.to_device(queue, + np.array( + range(nboxes - n_active_boxes, nboxes), + dtype=self.box_id_dtype)) + assert len(self.active_boxes) == 1 << ((nlevels - 1) * dim) + self.register_fields({"active_boxes": self.active_boxes}) + + # FIXME: map bool in pyopencl + # pyopencl/compyte/dtypes.py", line 107, in dtype_to_ctype + # raise ValueError("unable to map dtype '%s'" % dtype) + # ValueError: unable to map dtype 'bool' + self.box_is_active = cl.array.zeros(queue, nboxes, np.int8) + self.box_is_active[nboxes - n_active_boxes:] = 1 + self.register_fields({"box_is_active": self.box_is_active}) + + self.level_boxes = make_obj_array([ + cl.array.zeros(queue, 1 << (dim * l), self.box_id_dtype) + for l in range(nlevels)]) + ibox = 0 + for l in range(nlevels): + for b in range(len(self.level_boxes[l])): + self.level_boxes[l][b] = ibox + ibox += 1 + self.register_fields({"level_boxes": self.level_boxes}) + + self.box_parent_ids = cl.array.zeros(queue, nboxes, self.box_id_dtype) + self.box_child_ids = cl.array.zeros(queue, (1 << dim, nboxes), + self.box_id_dtype) + for l in range(nlevels): + if l == 0: # noqa: E741 + self.box_parent_ids[0] = 0 + else: + multi_index_bases = tuple( + 1 << ((dim - 1 - d) * (l-1)) for d in range(dim)) + for ilb, multi_ind in zip(range(len(self.level_boxes[l])), + product(range(1 << l), repeat=dim)): + ibox = self.box_id_dtype(self.level_boxes[l][ilb].get()) + parent_multi_ind = tuple(ind // 2 for ind in multi_ind) + parent_level_id = np.sum([ind * base for ind, base + in zip(parent_multi_ind, multi_index_bases)]) + self.box_parent_ids[ibox] = self.level_boxes[ + l-1][parent_level_id] + + child_multi_index_bases = tuple( + 1 << (dim - 1 - d) for d in range(dim)) + child_multi_ind = tuple(ind - pind * 2 for ind, pind + in zip(multi_ind, parent_multi_ind)) + child_id = np.sum([ind * base for ind, base + in zip(child_multi_ind, child_multi_index_bases)]) + self.box_child_ids[child_id][self.box_id_dtype( + self.level_boxes[l-1][parent_level_id].get()) + ] = ibox + self.register_fields({ + "box_parent_ids": self.box_parent_ids, + "box_child_ids": self.box_child_ids}) + + self.adaptivity_flags = cl.array.zeros(queue, + self.n_active_boxes, + adaptivity_flags_enum.dtype) + self.register_fields({"adaptivity_flags": self.adaptivity_flags}) + + n_active_neighbors_host = np.zeros(n_active_boxes, dtype=self.size_t) + nbpatch = [np.array(patch) + for patch in product(range(-1, 2), repeat=dim) + if sum(np.abs(patch)) > 0] + assert len(nbpatch) == 3 ** dim - 1 + maxid = (1 << (nlevels - 1)) - 1 + iabox = 0 + for multiabox in product(range(1 << (nlevels - 1)), repeat=dim): + if ( + min(multiabox) > 0 + and max(multiabox) < maxid + ): + n_active_neighbors_host[iabox] = 3 ** dim - 1 + else: + n_active_neighbors_host[iabox] = sum(1 for patch in nbpatch + if ( + np.min(np.array(multiabox) + patch) >= 0 + and np.max(np.array(multiabox) + patch) <= maxid + )) + iabox += 1 + self.n_active_neighbors = cl.array.to_device(queue, + n_active_neighbors_host) + self.register_fields({"n_active_neighbors": self.n_active_neighbors}) + + ind_bases = np.array([1 << (nlevels - 1) * (dim - d - 1) + for d in range(dim)]) + anb_pre = [] + iabox = 0 + for multiabox in product(range(1 << (nlevels - 1)), repeat=dim): + nbid = 0 + anb_pre.append(np.zeros( + n_active_neighbors_host[iabox], dtype=self.box_id_dtype)) + for patch in nbpatch: + tmp = np.array(multiabox) + patch + if np.min(tmp) >= 0 and np.max(tmp) <= maxid: + anb_pre[-1][nbid] = self.box_id_dtype( + nboxes - n_active_boxes + + np.sum(tmp * ind_bases)) + nbid += 1 + iabox += 1 + self.active_neighboring_boxes = make_obj_array([ + cl.array.to_device(queue, anb) for anb in anb_pre]) + self.register_fields({ + "active_neighboring_boxes": self.active_neighboring_boxes}) + + @property + def dimensions(self): + return len(self.box_centers) + + @property + def nboxes(self): + return len(self.box_levels) + + @property + def nlevels(self): + return len(self.level_boxes) + + @property + def n_active_boxes(self): + return len(self.active_boxes) + + def plot(self, **kwargs): + from boxtree.visualization import BoxTreePlotter + plotter = BoxTreePlotter(self) + plotter.draw_tree(**kwargs) + plotter.set_bounding_box() + + def get_box_extent(self, ibox): + if isinstance(ibox, cl.array.Array): + ibox = self.box_id_dtype(ibox.get()) + lev = self.box_level_dtype(self.box_levels[ibox].get()) + else: + lev = self.box_level_dtype(self.box_levels[ibox]) + box_size = self.root_extent / (1 << lev) + extent_low = np.zeros(self.dimensions, self.coord_dtype) + for d in range(self.dimensions): + if isinstance(self.box_centers[0], cl.array.Array): + extent_low[d] = self.box_centers[ + d, ibox].get() - 0.5 * box_size + else: + extent_low[d] = self.box_centers[d, ibox] - 0.5 * box_size + + extent_high = extent_low + box_size + return extent_low, extent_high + +# }}} End Data structure for a tree of boxes + +# {{{ Discretize a BoxTree using quadrature + + +class QuadratureOnBoxTree(object): + """Tensor-product quadrature rules applied to each active box. + + This class only holds data on template quadrature rules, while it + responds to queries for: + - quadrature nodes + - quadrature weights + - active box centers + - active box measures (areas/volumes) + + The information is generated on the fly and is responsive to + changes made to the underlying BoxTree object. + + .. attribute:: quadrature_formula + + The 1D template quadrature formula defined on [-1, 1]. + + .. attribute:: boxtree + + The underlying BoxTree object, assumed to reside on the device, + i.e., its get() method was not called. + """ + + def __init__(self, boxtree, quadrature_formula=None): + """If no quadrature_formula is supplied, the trivial one is used + sum[ f(centers) * measures ] + """ + if not isinstance(boxtree, BoxTree): + raise RuntimeError("Invalid boxtree") + + self.boxtree = boxtree + + if quadrature_formula is None: + from modepy import LegendreGaussQuadrature + self.quadrature_formula = LegendreGaussQuadrature(0) + else: + self.quadrature_formula = quadrature_formula + + def get_q_points(self, queue): + """Return quadrature points. + """ + q_nodes = cl.array.to_device(queue, self.quadrature_formula.nodes) + + def tensor_q_node_code(dim): + if dim == 1: + return "q_nodes[iq0]" + elif dim == 2: + return "if(iaxis==0, q_nodes[iq0], q_nodes[iq1])" + elif dim == 3: + return """if(iaxis==0, q_nodes[iq0], + if(iaxis==1, q_nodes[iq1], q_nodes[iq2]) + ) + """ + + def get_knl(dim): + quad_vars = ["iq%d" % i for i in range(dim)] + knl = lp.make_kernel( + "{[iabox, iaxis, QUAD_VARS]".replace( + "QUAD_VARS", ', '.join(quad_vars)) + + ": 0<=iabox box_id = active_boxes[iabox] + <> box_level = box_levels[box_id] + <> box_relative_extent = root_extent / 2 ** (box_level+1) + for iaxis, QUAD_VARS + result[iaxis, iabox, QUAD_VARS] = box_centers[ + iaxis, box_id] + ( + GET_TENSOR_Q_NODE_IAXIS * box_relative_extent) + end + end + """ + .replace("GET_TENSOR_Q_NODE_IAXIS", tensor_q_node_code(dim)) + .replace("QUAD_VARS", ", ".join(quad_vars)), + [ + lp.GlobalArg("result", self.boxtree.coord_dtype, + shape=lp.auto), + lp.GlobalArg("active_boxes", self.boxtree.box_id_dtype, + shape=lp.auto), + lp.GlobalArg("box_levels", self.boxtree.box_level_dtype, + shape="nboxes"), + lp.GlobalArg("q_nodes", self.boxtree.coord_dtype, + shape=lp.auto), + lp.GlobalArg("box_centers", self.boxtree.coord_dtype, + shape="(DIM, nboxes)".replace("DIM", str(dim))), + lp.ValueArg("root_extent", self.boxtree.coord_dtype), + lp.ValueArg("n_active_boxes", self.boxtree.size_t), + lp.ValueArg("n_q_nodes", self.boxtree.size_t), + lp.ValueArg("nboxes", self.boxtree.size_t), + ], + name="get_active_box_quad_points") + + knl = lp.split_iname(knl, + "iabox", 128, outer_tag="g.0", inner_tag="l.0") + + return knl + + evt, result = get_knl(self.boxtree.dimensions)(queue, + active_boxes=self.boxtree.active_boxes, + box_levels=self.boxtree.box_levels, + q_nodes=q_nodes, + box_centers=self.boxtree.box_centers, + root_extent=self.boxtree.root_extent, + n_active_boxes=self.boxtree.n_active_boxes, + n_q_nodes=len(q_nodes), + nboxes=self.boxtree.nboxes) + + assert len(result) == 1 + # TODO: If the ordering is not the same as dealii, + # reverse the order of quad vars + return make_obj_array([cpnt.ravel() for cpnt in result[0]]) + + def get_q_weights(self, queue): + """Return quadrature weights. + """ + q_weights = cl.array.to_device(queue, self.quadrature_formula.weights) + + def get_knl(dim): + quad_vars = ["iq%d" % i for i in range(dim)] + knl = lp.make_kernel( + "{[iabox, QUAD_VARS]:".replace( + "QUAD_VARS", ', '.join(quad_vars)) + + " 0<=iabox box_id = active_boxes[iabox] + <> box_level = box_levels[box_id] + <> box_relative_measure = root_measure / ( + 2 ** ((box_level+1) * DIM)) + for QUAD_VARS + result[iabox, QUAD_VARS] = (QUAD_WEIGHT + * box_relative_measure) + end + end + """.replace("DIM", str(dim)) + .replace("QUAD_VARS", ", ".join(quad_vars)) + .replace("QUAD_WEIGHT", " * ".join( + ["q_weights[" + qvar + "]" for qvar in quad_vars])), + [ + lp.GlobalArg("result", self.boxtree.coord_dtype, + shape=lp.auto), + lp.GlobalArg("active_boxes", self.boxtree.box_id_dtype, + shape=lp.auto), + lp.GlobalArg("box_levels", self.boxtree.box_level_dtype, + shape="nboxes"), + lp.GlobalArg("q_weights", self.boxtree.coord_dtype, + shape=lp.auto), + lp.ValueArg("root_measure", self.boxtree.coord_dtype), + lp.ValueArg("n_active_boxes", self.boxtree.size_t), + lp.ValueArg("n_q_nodes", self.boxtree.size_t), + lp.ValueArg("nboxes", self.boxtree.size_t), + ], + name="get_active_box_quad_weights") + + knl = lp.split_iname(knl, + "iabox", 128, outer_tag="g.0", inner_tag="l.0") + + return knl + + evt, result = get_knl(self.boxtree.dimensions)(queue, + active_boxes=self.boxtree.active_boxes, + box_levels=self.boxtree.box_levels, + q_weights=q_weights, + root_measure=(self.boxtree.root_extent + ** self.boxtree.dimensions), + n_active_boxes=self.boxtree.n_active_boxes, + n_q_nodes=len(q_weights), + nboxes=self.boxtree.nboxes) + + assert len(result) == 1 + # TODO: If the ordering is not the same as dealii, + # reverse the order of quad vars + return result[0].ravel() + + def get_cell_centers(self, queue): + """Return centers of active boxes. + """ + def get_knl(): + knl = lp.make_kernel( # noqa: E131 + "{[iabox, iaxis]: 0<=iabox box_id = active_boxes[iabox] + result[iaxis, iabox] = box_centers[iaxis, box_id] + end + """, + [ + lp.GlobalArg("result", self.boxtree.coord_dtype, + shape=lp.auto), + lp.GlobalArg("active_boxes", self.boxtree.box_id_dtype, + shape=lp.auto), + lp.GlobalArg("box_centers", self.boxtree.coord_dtype, + shape="(dims, nboxes)"), + lp.ValueArg("n_active_boxes", self.boxtree.size_t), + lp.ValueArg("dims", self.boxtree.size_t), + lp.ValueArg("nboxes", self.boxtree.size_t), + ], + name="get_active_box_centers") + + knl = lp.split_iname(knl, + "iabox", 128, outer_tag="g.0", inner_tag="l.0") + + return knl + + evt, result = get_knl()(queue, + active_boxes=self.boxtree.active_boxes, + box_centers=self.boxtree.box_centers, + n_active_boxes=self.boxtree.n_active_boxes, + dims=self.boxtree.dimensions, + nboxes=self.boxtree.nboxes) + + assert len(result) == 1 + return make_obj_array([cpnt.ravel() for cpnt in result[0]]) + + def get_cell_measures(self, queue): + """Return measures of active boxes. + This method does not call BoxTree.get_box_extent() for performance + reasons. + """ + def get_knl(): + knl = lp.make_kernel( + "{[iabox]: 0<=iabox box_id = active_boxes[iabox] + <> box_level = box_levels[box_id] + result[iabox] = root_measure / (2 ** (box_level * dims)) + end + """, + [ + lp.GlobalArg("result", self.boxtree.coord_dtype, + shape=lp.auto), + lp.GlobalArg("active_boxes", self.boxtree.box_id_dtype, + shape=lp.auto), + lp.GlobalArg("box_levels", self.boxtree.box_level_dtype, + shape="nboxes"), + lp.ValueArg("root_measure", self.boxtree.coord_dtype), + lp.ValueArg("n_active_boxes", self.boxtree.size_t), + lp.ValueArg("dims", self.boxtree.size_t), + lp.ValueArg("nboxes", self.boxtree.size_t), + ], + name="get_active_box_measures") + + knl = lp.split_iname(knl, + "iabox", 128, outer_tag="g.0", inner_tag="l.0") + + return knl + + evt, result = get_knl()(queue, + active_boxes=self.boxtree.active_boxes, + box_levels=self.boxtree.box_levels, + root_measure=( + self.boxtree.root_extent ** self.boxtree.dimensions), + n_active_boxes=self.boxtree.n_active_boxes, + dims=self.boxtree.dimensions, + nboxes=self.boxtree.nboxes) + + assert len(result) == 1 + return result[0] + + +# }}} End Discretize a BoxTree using quadrature diff --git a/boxtree/visualization.py b/boxtree/visualization.py index e9671e07deb961eb9cd3cd9bf53447d5dcf69a67..927142503924f372c97269e2d74ecc758739c022 100644 --- a/boxtree/visualization.py +++ b/boxtree/visualization.py @@ -172,9 +172,105 @@ class TreePlotter: # }}} +# {{{ box tree plotting + + +class BoxTreePlotter(TreePlotter): + """Assumes that the tree has data living on the host. + See :meth:`boxtree.BoxTree.get`. + """ + + def __init__(self, tree): + self.tree = tree + + def draw_tree(self, **kwargs): + if self.tree.dimensions != 2: + raise NotImplementedError("can only plot 2D trees for now") + + fill = kwargs.pop("fill", False) + edgecolor = kwargs.pop("edgecolor", "black") + kwargs["fill"] = fill + kwargs["edgecolor"] = edgecolor + + for iabox in range(self.tree.n_active_boxes): + ibox = self.tree.box_id_dtype(self.tree.active_boxes[iabox]) + self.draw_box(ibox, **kwargs) + + def set_bounding_box(self): + import matplotlib.pyplot as pt + root_center = np.array([self.tree.box_centers[d, 0] for d + in range(self.tree.dimensions)]) + root_extent = self.tree.root_extent + pt.xlim(root_center[0] - root_extent / 2, + root_center[0] + root_extent / 2) + pt.ylim(root_center[1] - root_extent / 2, + root_center[1] + root_extent / 2) + + pt.gca().set_aspect("equal") + + def draw_box_numbers(self): + import matplotlib.pyplot as pt + + tree = self.tree + + for iabox in range(tree.n_active_boxes): + ibox = tree.box_id_dtype(tree.active_boxes[iabox]) + x, y = tree.box_centers[:, ibox] + lev = int(tree.box_levels[ibox]) + pt.text(x, y, str(ibox), fontsize=20*1.15**(-lev), + ha="center", va="center", + bbox=dict(facecolor='white', alpha=0.5, lw=0)) + + def get_tikz_for_tree(self): + if self.tree.dimensions != 2: + raise NotImplementedError("can only plot 2D trees for now") + + lines = [] + + lines.append(r"\def\nboxes{%d}" % self.tree.nboxes) + lines.append(r"\def\lastboxnr{%d}" % (self.tree.nboxes-1)) + + for iabox in range(self.tree.n_active_boxes): + ibox = self.tree.box_id_dtype(self.tree.active_boxes[iabox]) + el, eh = self.tree.get_box_extent(ibox) + + c = self.tree.box_centers[:, ibox] + + lines.append( + r"\coordinate (boxl%d) at (%r, %r);" + % (ibox, float(el[0]), float(el[1]))) + lines.append( + r"\coordinate (boxh%d) at (%r, %r);" + % (ibox, float(eh[0]), float(eh[1]))) + lines.append( + r"\coordinate (boxc%d) at (%r, %r);" + % (ibox, float(c[0]), float(c[1]))) + lines.append( + r"\def\boxsize%s{%r}" + % (int_to_roman(ibox), float(eh[0]-el[0]))) + lines.append( + r"\def\boxlevel%s{%r}" + % (int_to_roman(ibox), self.tree.box_levels[ibox])) + + lines.append( + r"\def\boxpath#1{(boxl#1) rectangle (boxh#1)}") + lines.append( + r"\def\drawboxes{" + r"\foreach \ibox in {0,...,\lastboxnr}{" + r"\draw \boxpath{\ibox};" + r"}}") + lines.append( + r"\def\drawboxnrs{" + r"\foreach \ibox in {0,...,\lastboxnr}{" + r"\node [font=\tiny] at (boxc\ibox) {\ibox};" + r"}}") + return "\n".join(lines) + +# }}} # {{{ traversal plotting + def _draw_box_list(tree_plotter, ibox, starts, lists, key_to_box=None, **kwargs): default_facecolor = "blue" diff --git a/examples/interactive_boxtree_demo.py b/examples/interactive_boxtree_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..1a299058bedf3a5b9a83310fc83df74ae99a0185 --- /dev/null +++ b/examples/interactive_boxtree_demo.py @@ -0,0 +1,58 @@ +import pyopencl as cl +import boxtree.tree_interactive_build + +import matplotlib.pyplot as pt +import numpy as np + +ctx = cl.create_some_context() +queue = cl.CommandQueue(ctx) + +tree = boxtree.tree_interactive_build.BoxTree() +tree.generate_uniform_boxtree(queue, nlevels=3, root_extent=2, + # root_vertex=(0,0,0)) + root_vertex=(0,0)) + +# the default quad formula uses cell centers and cell measures +from modepy import GaussLegendreQuadrature +n_q_points = 5 +quadrature_formula = GaussLegendreQuadrature(n_q_points - 1) +print(quadrature_formula.nodes, quadrature_formula.weights) +quad = boxtree.tree_interactive_build.QuadratureOnBoxTree(tree, + quadrature_formula) +cell_centers = quad.get_cell_centers(queue) +cell_measures = quad.get_cell_measures(queue) +q_points = quad.get_q_points(queue) +q_weights = quad.get_q_weights(queue) + +# print(q_points) +# print(q_weights, np.sum(q_weights.get())) + +# print(q_points - cell_centers) +# print(q_weights - cell_measures) + +# call get() before plotting +from boxtree.visualization import BoxTreePlotter +plt = BoxTreePlotter(tree.get(queue)) +plt.draw_tree() +plt.set_bounding_box() +plt.draw_box_numbers() + +qx = q_points[0].get(queue) +qy = q_points[1].get(queue) + +pt.plot(qx, qy, '*') + +pt.tick_params( + axis='x', # changes apply to the x-axis + which='both', # both major and minor ticks are affected + bottom='off', # ticks along the bottom edge are off + top='off', # ticks along the top edge are off + labelbottom='off') +pt.tick_params( + axis='y', + which='both', + left='off', + top='off', + labelleft='off') + +pt.show() diff --git a/setup.py b/setup.py index 0882bd2475b8e5d58ccaf9d8d3478fec9a124b13..d1cb8b5ae5f168df0e80532964f851cf08b47541 100644 --- a/setup.py +++ b/setup.py @@ -44,6 +44,7 @@ def main(): "pyopencl>=2018.2.2", "Mako>=0.7.3", "pytest>=2.3", + "modepy>=2016.1.2", "cgen>=2013.1.2", "six", ])