Very drafty work towards new tree build

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

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