diff --git a/boxtree/traversal.py b/boxtree/traversal.py index 74af3319..4691d583 100644 --- a/boxtree/traversal.py +++ b/boxtree/traversal.py @@ -1792,10 +1792,6 @@ def get_kernel_info(self, dimensions, particle_id_dtype, box_id_dtype, "same_level_non_well_sep_boxes_starts"), VectorArg(box_id_dtype, "same_level_non_well_sep_boxes_lists"), - VectorArg(coord_dtype, "box_target_bounding_box_min", - with_offset=False), - VectorArg(coord_dtype, "box_target_bounding_box_max", - with_offset=False), VectorArg(particle_id_dtype, "box_source_counts_cumul"), ScalarArg(particle_id_dtype, "from_sep_smaller_min_nsources_cumul"), @@ -2034,9 +2030,7 @@ def extract_level_start_box_nrs(box_list, wait_for): tree.stick_out_factor, target_boxes, same_level_non_well_sep_boxes.starts, same_level_non_well_sep_boxes.lists, - tree.box_target_bounding_box_min.data, - tree.box_target_bounding_box_max.data, - tree.box_source_counts_cumul, + tree.box_source_counts_cumul if with_extent else None, _from_sep_smaller_min_nsources_cumul, ) @@ -2192,4 +2186,4 @@ def extract_level_start_box_nrs(box_list, wait_for): # }}} -# vim: filetype=pyopencl:fdm=marker +# vim: fdm=marker diff --git a/boxtree/tree.py b/boxtree/tree.py index 68f06add..58d5e0db 100644 --- a/boxtree/tree.py +++ b/boxtree/tree.py @@ -28,11 +28,12 @@ .. autoclass:: Tree +.. currentmodule:: boxtree.tree +.. autoclass:: TreeOfBoxes + Tree with linked point sources ------------------------------ -.. currentmodule:: boxtree.tree - .. autoclass:: TreeWithLinkedPointSources .. autofunction:: link_point_sources @@ -83,9 +84,11 @@ import pyopencl as cl import numpy as np +import numpy.typing as npt from boxtree.tools import DeviceDataRecord from cgen import Enum from pytools import memoize_method +from dataclasses import dataclass import logging logger = logging.getLogger(__name__) @@ -111,9 +114,106 @@ class box_flags_enum(Enum): # noqa # }}} +# {{{ tree of boxes + +@dataclass +class TreeOfBoxes: + """A quad/octree tree of abstract boxes. Lightweight trees handled with + :mod:`numpy`, intended for mesh adaptivity. + + .. attribute:: dimensions + + .. attribute:: nlevels + + .. attribute:: nboxes + + .. attribute:: root_extent + + (Scalar) extent of the root box. + + .. attribute:: box_centers + + dim-by-nboxes :mod:`numpy` array of the centers of the boxes. + + .. attribute:: box_parent_ids + + :mod:`numpy` vector of parent box ids. + + .. attribute:: box_child_ids + + (2**dim)-by-nboxes :mod:`numpy` array of children box ids. + + .. attribute:: box_levels + + :mod:`numpy` vector of box levels in non-decreasing order. + + .. automethod:: __init__ + + .. automethod:: get_leaf_flags + + .. automethod:: leaf_boxes + """ + box_centers: npt.NDArray + root_extent: npt.NDArray + box_parent_ids: npt.NDArray + box_child_ids: npt.NDArray + box_levels: npt.NDArray + + def __post_init__(self): + if isinstance(self.box_centers, cl.array.Array): + bcenters = self.box_centers.get() + else: + bcenters = self.box_centers + lows = bcenters[:, 0] - 0.5*self.root_extent + highs = lows + self.root_extent + self.bounding_box = [lows, highs] + self.dim = self.box_centers.shape[0] + + # {{{ dummy interface for TreePlotter + + @property + def dimensions(self): + return self.dim + + def get_box_size(self, ibox): + lev = self.box_levels[ibox] + box_size = self.root_extent * 0.5**lev + return box_size + + def get_box_extent(self, ibox): + box_size = self.get_box_size(ibox) + extent_low = self.box_centers[:, ibox] - 0.5*box_size + extent_high = extent_low + box_size + return extent_low, extent_high + + # }}} End dummy interface for TreePlotter + + @property + def nboxes(self): + return self.box_centers.shape[1] + + @property + def nlevels(self): + # level starts from 0 + if isinstance(self.box_levels, cl.array.Array): + return int(max(self.box_levels).get()) + 1 + else: + return max(self.box_levels) + 1 + + def get_leaf_flags(self): + # box_id -> whether the box is leaf + return np.all(self.box_child_ids == 0, axis=0) + + def leaf_boxes(self): + boxes = np.arange(self.nboxes) + return boxes[self.get_leaf_flags()] + +# }}} + + # {{{ tree data structure -class Tree(DeviceDataRecord): +class Tree(DeviceDataRecord, TreeOfBoxes): r"""A quad/octree consisting of particles sorted into a hierarchy of boxes. Optionally, particles may be designated 'sources' and 'targets'. They may also be assigned radii which restrict the minimum size of the box @@ -400,7 +500,6 @@ class Tree(DeviceDataRecord): .. automethod:: get """ - @property def dimensions(self): return len(self.sources) @@ -440,6 +539,12 @@ def get_box_extent(self, ibox): extent_high = extent_low + box_size return extent_low, extent_high + def get_leaf_flags(self): + raise NotImplementedError + + def leaf_boxes(self): + raise NotImplementedError + # {{{ debugging aids # these assume numpy arrays (i.e. post-.get()), for now diff --git a/boxtree/tree_build.py b/boxtree/tree_build.py index cf70c573..9d147d68 100644 --- a/boxtree/tree_build.py +++ b/boxtree/tree_build.py @@ -5,6 +5,16 @@ -------------- .. autoclass:: TreeBuilder + +Manipulating Trees of Boxes +--------------------------- + +.. currentmodule:: boxtree.tree_build + +.. autofunction:: refined +.. autofunction:: uniformly_refined +.. autofunction:: coarsened +.. autofunction:: refined_and_coarsened """ __copyright__ = "Copyright (C) 2012 Andreas Kloeckner" @@ -31,11 +41,13 @@ import numpy as np +import numpy.typing as npt +from typing import Union, TYPE_CHECKING, List from pytools import memoize_method import pyopencl as cl import pyopencl.array # noqa from functools import partial -from boxtree.tree import Tree +from boxtree.tree import Tree, TreeOfBoxes from pytools import ProcessLogger, DebugProcessLogger import logging @@ -46,6 +58,272 @@ class MaxLevelsExceeded(RuntimeError): pass +# {{{ utils for tree of boxes + +def _resized_array(arr: npt.NDArray, new_size: int) -> npt.NDArray: + """Return a resized copy of the array. The new_size is a scalar which is + applied to the last dimension. + """ + old_size = arr.shape[-1] + prefix = (slice(None), ) * (arr.ndim - 1) + if old_size >= new_size: + return arr[prefix + (slice(new_size), )].copy() + else: + new_shape = list(arr.shape) + new_shape[-1] = new_size + new_arr = np.zeros(new_shape, arr.dtype) + new_arr[prefix + (slice(old_size), )] = arr + return new_arr + + +def _vec_of_signs(dim: int, i: int) -> npt.NDArray: + """The sign vector is obtained by converting i to a dim-bit binary. + """ + # e.g. bin(10) = '0b1010' + binary_digits = [int(bd) for bd in bin(i)[2:]] + n = len(binary_digits) + assert n <= dim + return np.array([0]*(dim-n) + binary_digits) * 2 - 1 + + +# {{{ refine/coarsen a tree of boxes + +def refined(tob: TreeOfBoxes, refine_flags: npt.NDArray) -> TreeOfBoxes: + """Make a refined copy of `tob` where boxes flagged with `refine_flags` are + refined. + """ + return refined_and_coarsened(tob, refine_flags, None) + + +def uniformly_refined(tob: TreeOfBoxes) -> TreeOfBoxes: + """Make a uniformly refined copy of `tob`. + """ + refine_flags = np.zeros(tob.nboxes, bool) + refine_flags[tob.get_leaf_flags()] = 1 + return refined(tob, refine_flags) + + +def coarsened(tob: TreeOfBoxes, coarsen_flags: npt.NDArray, + error_on_ignored_flags: bool = True + ) -> TreeOfBoxes: + """Make a coarsened copy of `tob` where boxes flagged with `coarsen_flags` + are coarsened. + """ + return refined_and_coarsened( + tob, None, coarsen_flags, error_on_ignored_flags) + + +def _apply_refine_flags_without_sorting(refine_flags, tob): + if refine_flags[~tob.get_leaf_flags()].any(): + raise ValueError("attempting to split non-leaf") + + refine_parents, = np.where(refine_flags) + if len(refine_parents) == 0: + return tob + + nchildren = 2**tob.dim + n_new_boxes = len(refine_parents) * nchildren + nboxes_new = tob.nboxes + n_new_boxes + + child_box_starts = ( + tob.nboxes + + nchildren * np.arange(len(refine_parents))) + + refine_parents_per_child = np.empty( + (nchildren, len(refine_parents)), np.intp) + refine_parents_per_child[:] = refine_parents.reshape(-1) + refine_parents_per_child = refine_parents_per_child.reshape(-1) + + box_parents = _resized_array(tob.box_parent_ids, nboxes_new) + box_centers = _resized_array(tob.box_centers, nboxes_new) + box_children = _resized_array(tob.box_child_ids, nboxes_new) + box_levels = _resized_array(tob.box_levels, nboxes_new) + + # new boxes are appended at the end, so applying coarsen_flags wrt the + # original tree is still meaningful after this + box_parents[tob.nboxes:] = refine_parents_per_child + box_levels[tob.nboxes:] = tob.box_levels[box_parents[tob.nboxes:]] + 1 + box_children[:, refine_parents] = ( + child_box_starts + np.arange(nchildren).reshape(-1, 1)) + + for i in range(2**tob.dim): + children_i = box_children[i, refine_parents] + offsets = ( + tob.root_extent * _vec_of_signs(tob.dim, i).reshape(-1, 1) + * (1/2**(1+box_levels[children_i]))) + box_centers[:, children_i] = ( + box_centers[:, refine_parents] + offsets) + + return TreeOfBoxes( + box_centers=box_centers, root_extent=tob.root_extent, + box_parent_ids=box_parents, box_child_ids=box_children, + box_levels=box_levels) + + +def _apply_coarsen_flags(coarsen_flags, tob, error_on_ignored_flags=True): + if coarsen_flags[~tob.get_leaf_flags()].any(): + raise ValueError("attempting to coarsen non-leaf") + coarsen_sources, = np.where(coarsen_flags) + if not coarsen_sources: + return tob + + coarsen_parents = tob.box_parent_ids[coarsen_sources] + coarsen_peers = tob.box_child_ids[:, coarsen_parents].reshape(-1) + coarsen_peer_is_leaf = tob.get_leaf_flags()[coarsen_peers] + coarsen_exec_flags = np.all(coarsen_peer_is_leaf, axis=0) + + # when a leaf box marked for coarsening has non-leaf peers + coarsen_flags_ignored = (coarsen_exec_flags != coarsen_flags) + if np.any(coarsen_flags_ignored): + msg = (f"{np.sum(coarsen_flags_ignored)} out of " + f"{np.sum(coarsen_flags)} coarsening flags ignored " + "to prevent removing non-leaf boxes") + if error_on_ignored_flags: + raise RuntimeError(msg) + else: + import warnings + warnings.warn(msg) + + # deleted boxes are marked as: + # level = inf + # parent = -1 + coarsen_parents = coarsen_parents[coarsen_exec_flags] + coarsen_peers = coarsen_peers[:, coarsen_exec_flags] + box_parents = tob.box_parent_ids.copy() + box_parents[coarsen_peers] = -1 + box_children = tob.box_child_ids.copy() + box_children[:, coarsen_parents] = 0 + box_levels = tob.box_levels.copy() + box_levels[coarsen_peers] = np.inf + + from dataclasses import replace + return replace(tob, box_parent_ids=box_parents, + box_child_ids=box_children, box_levels=box_levels) + + +def _sort_boxes_by_level(tob, queue=None): + if np.any(np.diff(tob.box_levels) < 0): + # reorder boxes to into non-decreasing levels + neworder = np.argsort(tob.box_levels) + box_centers = tob.box_centers[:, neworder] + box_parent_ids = tob.box_parent_ids[neworder] + box_child_ids = tob.box_child_ids[:, neworder] + box_levels = tob.box_levels[neworder] + tob = TreeOfBoxes( + box_centers=box_centers, root_extent=tob.root_extent, + box_parent_ids=box_parent_ids, box_child_ids=box_child_ids, + box_levels=box_levels) + return tob + + +def _sort_and_prune_deleted_boxes(tob): + tob = _sort_boxes_by_level(tob) + n_stale_boxes = np.sum(tob.box_levels == np.inf) + newn = tob.nboxes - n_stale_boxes + from dataclasses import replace + return replace(tob, box_parent_ids=tob.box_parent_ids[:newn], + box_child_ids=tob.box_child_ids[:, :newn], + box_levels=tob.box_levels[:newn], + box_centers=tob.box_centers[:, :newn]) + + +def refined_and_coarsened(tob: TreeOfBoxes, + refine_flags: Union[None, npt.NDArray] = None, + coarsen_flags: Union[None, npt.NDArray] = None, + error_on_ignored_flags: bool = True + ) -> TreeOfBoxes: + """Make a refined/coarsened copy. When children of the same parent box + are marked differently, the refinement flag takes priority. + + Both refinement and coarsening flags can only be set of leaves. + To prevent drastic mesh change, coarsening is only executed when a leaf + box is marked for coarsening, and its parent's children are all leaf + boxes (so that change in the number of boxes is bounded per box flagged). + + :arg refine_flags: a boolean array of size `nboxes`. + :arg coarsen_flags: a boolean array of size `nboxes`. + :arg error_on_ignored_flags: if true, an exception is raised when enforcing + level restriction requires ignoring some coarsening flags. + :returns: a processed copy of the tree. + """ + if refine_flags is None: + refine_flags = np.zeros(tob.nboxes, bool) + if coarsen_flags is None: + coarsen_flags = np.zeros(tob.nboxes, bool) + + if (refine_flags & coarsen_flags).any(): + raise ValueError("some boxes are simultaneously marked " + "to refine and coarsen") + + tob = _apply_refine_flags_without_sorting(refine_flags, tob) + coarsen_flags = _resized_array(coarsen_flags, tob.nboxes) + tob = _apply_coarsen_flags(coarsen_flags, tob, error_on_ignored_flags) + return _sort_and_prune_deleted_boxes(tob) + +# }}} End refine/coarsen a tree of boxes + + +def make_tob_root(dim: int, + bbox: Union[npt.NDArray, List[List[float]]]) -> TreeOfBoxes: + """ + Make the minimal tree of boxes. + """ + center = np.array( + [(bbox[0][iaxis] + bbox[1][iaxis]) * 0.5 for iaxis in range(dim)]) + extents = np.array( + [(bbox[1][iaxis] - bbox[0][iaxis]) for iaxis in range(dim)]) + from pytools import single_valued + parents = np.array([0], np.intp) + parents[0] = -1 # root has no parent + return TreeOfBoxes( + box_centers=center.reshape(dim, 1), + root_extent=single_valued(extents, np.allclose), + box_parent_ids=parents, + box_child_ids=np.array([0] * 2**dim, np.intp).reshape(2**dim, 1), + box_levels=np.array([0]), + ) + + +if TYPE_CHECKING: + from meshmode.mesh import Mesh + + +def make_mesh_from_leaves(tob: TreeOfBoxes) -> "Mesh": + """Make a :mod:`meshmode` mesh from the leaf boxes. + + :arg tob: a :class:`TreeOfBoxes`. + """ + lfboxes = tob.leaf_boxes() + lfcenters = tob.box_centers[:, lfboxes] + lflevels = tob.box_levels[lfboxes] + lfradii = tob.root_extent / 2 / (2**lflevels) + + # use tensor product nodes ordering + import modepy.nodes as nd + cell_nodes_1d = np.array([-1, 1]) + cell_nodes = nd.tensor_product_nodes(tob.dim, cell_nodes_1d) + + lfvertices = ( + np.repeat(lfcenters, 2**tob.dim, axis=1) + + np.repeat(lfradii, 2**tob.dim) * np.tile(cell_nodes, (1, len(lfboxes))) + ) + + # FIXME: purge redundant vertices + from meshmode.mesh import Mesh, TensorProductElementGroup + from meshmode.mesh.generation import make_group_from_vertices + + vertex_indices = np.arange( + len(lfboxes) * 2**tob.dim, dtype=np.int32).reshape([-1, 2**tob.dim]) + group = make_group_from_vertices( + lfvertices, vertex_indices, 1, + group_cls=TensorProductElementGroup, + unit_nodes=None) + + return Mesh(vertices=lfvertices, groups=[group]) + +# }}} + + class TreeBuilder: """ .. automethod:: __init__ @@ -1733,7 +2011,6 @@ def make_arange(start): return Tree( # If you change this, also change the documentation # of what's in the tree, above. - sources_are_targets=sources_are_targets, sources_have_extent=sources_have_extent, targets_have_extent=targets_have_extent, diff --git a/doc/Makefile b/doc/Makefile index c45814ac..bb66c0e8 100644 --- a/doc/Makefile +++ b/doc/Makefile @@ -2,7 +2,7 @@ # # You can set these variables from the command line. -SPHINXOPTS = +SPHINXOPTS = -n SPHINXBUILD = python $(shell which sphinx-build) PAPER = BUILDDIR = _build diff --git a/doc/conf.py b/doc/conf.py index 07572b6c..e16f166d 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -20,3 +20,7 @@ "https://documen.tician.de/pyopencl/": None, "https://documen.tician.de/pytential/": None, } + +nitpick_ignore_regex = [ + ["py:class", r"numpy.typing._generic_alias.ScalarType"], +] diff --git a/requirements.txt b/requirements.txt index bd38c8c1..d1084939 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,8 @@ git+https://github.com/inducer/pyopencl#egg=pyopencl git+https://github.com/inducer/pymbolic#egg=pymbolic git+https://github.com/inducer/islpy#egg=islpy git+https://github.com/inducer/loopy#egg=loopy +git+https://github.com/inducer/modepy#egg=modepy +git+https://github.com/inducer/meshmode#egg=meshmode git+https://github.com/inducer/pyfmmlib#egg=pyfmmlib git+https://github.com/inducer/arraycontext#egg=arraycontext diff --git a/test/test_tree_of_boxes.py b/test/test_tree_of_boxes.py new file mode 100644 index 00000000..e8093ae2 --- /dev/null +++ b/test/test_tree_of_boxes.py @@ -0,0 +1,285 @@ +__copyright__ = "Copyright (C) 2012 Andreas Kloeckner, 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 sys +import pytest + +import numpy as np +import pyopencl as cl + +from arraycontext import pytest_generate_tests_for_array_contexts +from boxtree.array_context import ( # noqa: F401 + PytestPyOpenCLArrayContextFactory, _acf) + +import logging +logger = logging.getLogger(__name__) + +pytest_generate_tests = pytest_generate_tests_for_array_contexts([ + PytestPyOpenCLArrayContextFactory, + ]) + + +def make_global_leaf_quadrature(actx, tob, order): + from meshmode.discretization.poly_element import \ + GaussLegendreTensorProductGroupFactory + group_factory = GaussLegendreTensorProductGroupFactory(order=order) + + from boxtree.tree_build import make_mesh_from_leaves + mesh = make_mesh_from_leaves(tob) + + if 0: + from meshmode.mesh import visualization as mvis + import matplotlib.pyplot as plt + mvis.draw_2d_mesh(mesh, + set_bounding_box=True, + draw_vertex_numbers=False, + draw_element_numbers=False) + plt.plot(tob.box_centers[0][tob.leaf_boxes()], + tob.box_centers[1][tob.leaf_boxes()], "rx") + plt.plot(mesh.vertices[0], mesh.vertices[1], "ro") + plt.show() + + from meshmode.discretization import Discretization + discr = Discretization(actx, mesh, group_factory) + + lflevels = tob.box_levels[tob.leaf_boxes()] + lfmeasures = (tob.root_extent / (2**lflevels))**tob.dim + + from arraycontext import flatten + weights = flatten(discr.quad_weights(), actx).with_queue(actx.queue) + jacobians = cl.array.to_device( + actx.queue, + np.repeat(lfmeasures/(2**tob.dim), discr.groups[0].nunit_dofs)) + q = weights * jacobians + + from pytools.obj_array import make_obj_array + nodes = discr.nodes() + x = make_obj_array([flatten(coords, actx).with_queue(actx.queue) + for coords in nodes]) + + return x, q + + +@pytest.mark.parametrize("dim", [1, 2, 3]) +@pytest.mark.parametrize("order", [1, 2, 3]) +@pytest.mark.parametrize("nlevels", [1, 4]) +def test_uniform_tree_of_boxes(ctx_factory, dim, order, nlevels): + from boxtree.tree_build import make_tob_root, uniformly_refined + lower_bounds = np.random.rand(dim) + radius = np.random.rand() + 0.1 + upper_bounds = lower_bounds + radius + tob = make_tob_root(dim=dim, bbox=[lower_bounds, upper_bounds]) + + for _ in range(nlevels - 1): + tob = uniformly_refined(tob) + + from arraycontext import PyOpenCLArrayContext + queue = cl.CommandQueue(ctx_factory()) + actx = PyOpenCLArrayContext(queue) + + x, q = make_global_leaf_quadrature(actx, tob, order) + + # integrates 1 exactly + assert np.isclose(sum(q.get()), radius**dim) + + +@pytest.mark.parametrize("dim", [1, 2, 3]) +@pytest.mark.parametrize("order", [1, 2, 3, 4]) +def test_uniform_tree_of_boxes_convergence(ctx_factory, dim, order): + from boxtree.tree_build import make_tob_root, uniformly_refined + radius = np.pi + lower_bounds = np.zeros(dim) - radius/2 + upper_bounds = lower_bounds + radius + tob = make_tob_root(dim=dim, bbox=[lower_bounds, upper_bounds]) + + min_level = 0 + max_level = 1 + + for _ in range(min_level): + tob = uniformly_refined(tob) + + # integrate cos(0.1*x + 0.2*y + 0.3*z + e) over [-pi/2, pi/2]**dim + qexact_table = { + 1: 20 * np.sin(np.pi/20) * np.cos(np.e), + 2: 50 * (np.sqrt(5) - 1) * np.sin(np.pi/20) * np.cos(np.e), + 3: 250/3 * (np.sqrt(10 - 2*np.sqrt(5)) - 2) * np.cos(np.e) + } + qexact = qexact_table[dim] + + from pytools.convergence import EOCRecorder + eoc_rec = EOCRecorder() + + from arraycontext import PyOpenCLArrayContext + queue = cl.CommandQueue(ctx_factory()) + actx = PyOpenCLArrayContext(queue) + + for _ in range(min_level, max_level + 1): + x, q = make_global_leaf_quadrature(actx, tob, order) + x, q = (np.array([xx.get() for xx in x]), q.get()) + + inner = np.ones_like(q) * np.e + for iaxis in range(dim): + inner += (iaxis + 1) * 0.1 * x[iaxis] + f = np.cos(inner) + qh = np.sum(f * q) + err = abs(qexact - qh) + + if err < 1e-14: + break # eoc will be off after hitting machine epsilon + + # under uniform refinement, last box is always leaf + eoc_rec.add_data_point(tob.get_box_size(-1), err) + tob = uniformly_refined(tob) + + if len(eoc_rec.history) > 1: + # Gauss quadrature is exact up to degree 2q+1 + eps = 0.01 + assert eoc_rec.order_estimate() >= 2*order + 2 - eps + else: + print(err) + assert err < 1e-14 + + +def test_tree_rebuild_with_particle(ctx_factory): + from boxtree.tree_build import make_tob_root, uniformly_refined + radius = np.pi + dim = 2 + nlevels = 3 + lower_bounds = np.zeros(dim) - radius/2 + upper_bounds = lower_bounds + radius + tob = make_tob_root(dim=dim, bbox=[lower_bounds, upper_bounds]) + + for _ in range(nlevels - 1): + tob = uniformly_refined(tob) + + from arraycontext import PyOpenCLArrayContext + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + actx = PyOpenCLArrayContext(queue) + + # FIXME: tree builder sometimes fails to reproduce + nqpts = 4 + x, q = make_global_leaf_quadrature(actx, tob, order=nqpts - 1) + + from boxtree.tree_build import TreeBuilder + tb = TreeBuilder(cl_ctx) + tree, evt = tb(queue, x, + max_particles_in_box=nqpts**dim, + bbox=np.array([ + [lower_bounds[i], upper_bounds[i]] + for i in range(dim)])) + + if 0: + from boxtree.visualization import TreePlotter + import matplotlib.pyplot as plt + tp = TreePlotter(tob) + tp.draw_tree() + tp.set_bounding_box() + plt.plot(x[0].get(), x[1].get(), ".") + plt.show() + + if 0: + from boxtree.visualization import TreePlotter + import matplotlib.pyplot as plt + tp = TreePlotter(tree.get(queue)) + tp.draw_tree() + tp.set_bounding_box() + plt.plot(x[0].get(), x[1].get(), ".") + plt.show() + + +@pytest.mark.skip(reason="wip") +def test_traversal_from_tob(ctx_factory): + from boxtree.tree_build import make_tob_root, uniformly_refined + radius = np.pi + dim = 2 + nlevels = 3 + lower_bounds = np.zeros(dim) - radius/2 + upper_bounds = lower_bounds + radius + tob = make_tob_root(dim=dim, bbox=[lower_bounds, upper_bounds]) + + for _ in range(nlevels): + tob = uniformly_refined(tob) + + ctx = ctx_factory() + queue = cl.CommandQueue(ctx) + + from boxtree.tree_build import _sort_boxes_by_level + tob = _sort_boxes_by_level(tob) + + def _to_device(tob, queue): + from dataclasses import replace + return replace( + tob, + box_centers=cl.array.to_device(queue, tob.box_centers), + box_parent_ids=cl.array.to_device(queue, tob.box_parent_ids), + box_child_ids=cl.array.to_device(queue, tob.box_child_ids), + box_levels=cl.array.to_device(queue, tob.box_levels)) + + blvlist = tob.box_levels.tolist() + level_start_box_nrs = np.array( + [blvlist.index(lv) for lv in range(tob.nlevels)]) + + # FIXME: fill in data + tob = _to_device(tob, queue) + tob.level_start_box_nrs = level_start_box_nrs + tob.level_start_box_nrs_dev = cl.array.to_device(queue, level_start_box_nrs) + + tob._is_pruned = True + tob.sources_have_extent = False + tob.particle_id_dtype = np.dtype(np.int32) + tob.box_id_dtype = np.dtype(np.int32) + tob.box_level_dtype = np.dtype(np.int32) + tob.coord_dtype = np.dtype(np.float64) + tob.sources_are_targets = True + tob.targets_have_extent = False + tob.extent_norm = "linf" + tob.aligned_nboxes = tob.nboxes + + # particle sizes + # FIXME + tob.stick_out_factor = 1e-12 + tob.box_target_bounding_box_min = None + tob.box_source_bounding_box_min = None + + from boxtree.traversal import FMMTraversalBuilder + from boxtree.tree import box_flags_enum + from functools import partial + empty = partial(cl.array.empty, queue, allocator=None) + tob.box_flags = empty(tob.nboxes, box_flags_enum.dtype) + + tg = FMMTraversalBuilder(ctx) + # FIXME + # trav, _ = tg(queue, tob) + + +# You can test individual routines by typing +# $ python test_tree.py 'test_routine(cl.create_some_context)' + +if __name__ == "__main__": + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: fdm=marker