|
| 1 | +__copyright__ = "Copyright (C) 2022 Alexandru Fikl" |
| 2 | + |
| 3 | +__license__ = """ |
| 4 | +Permission is hereby granted, free of charge, to any person obtaining a copy |
| 5 | +of this software and associated documentation files (the "Software"), to deal |
| 6 | +in the Software without restriction, including without limitation the rights |
| 7 | +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 8 | +copies of the Software, and to permit persons to whom the Software is |
| 9 | +furnished to do so, subject to the following conditions: |
| 10 | +
|
| 11 | +The above copyright notice and this permission notice shall be included in |
| 12 | +all copies or substantial portions of the Software. |
| 13 | +
|
| 14 | +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 15 | +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 16 | +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 17 | +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 18 | +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 19 | +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN |
| 20 | +THE SOFTWARE. |
| 21 | +""" |
| 22 | + |
| 23 | +from dataclasses import dataclass, replace |
| 24 | +from functools import singledispatch |
| 25 | +from typing import Optional |
| 26 | + |
| 27 | +import numpy as np |
| 28 | + |
| 29 | +from pytools import T, memoize_method |
| 30 | + |
| 31 | +from arraycontext import PyOpenCLArrayContext |
| 32 | +from boxtree.tree import Tree |
| 33 | +from pytential import sym, GeometryCollection |
| 34 | +from pytential.linalg.utils import IndexList, TargetAndSourceClusterList |
| 35 | + |
| 36 | +__doc__ = """ |
| 37 | +Clustering |
| 38 | +~~~~~~~~~~ |
| 39 | +
|
| 40 | +.. autoclass:: ClusterTreeLevel |
| 41 | +
|
| 42 | +.. autofunction:: cluster |
| 43 | +.. autofunction:: partition_by_nodes |
| 44 | +""" |
| 45 | + |
| 46 | +# FIXME: this is just an arbitrary value |
| 47 | +_DEFAULT_MAX_PARTICLES_IN_BOX = 32 |
| 48 | + |
| 49 | + |
| 50 | +# {{{ cluster tree |
| 51 | + |
| 52 | +@dataclass(frozen=True) |
| 53 | +class ClusterTreeLevel: |
| 54 | + """ |
| 55 | + .. attribute:: level |
| 56 | +
|
| 57 | + Current level that is represented. |
| 58 | +
|
| 59 | + .. attribute:: nlevels |
| 60 | +
|
| 61 | + Total number of levels in the tree. |
| 62 | +
|
| 63 | + .. attribute:: nclusters |
| 64 | +
|
| 65 | + Number of clusters on the current level. |
| 66 | +
|
| 67 | + .. attribute:: partition_box_ids |
| 68 | +
|
| 69 | + Box IDs on the current level. |
| 70 | +
|
| 71 | + .. attribute:: partition_parent_ids |
| 72 | +
|
| 73 | + Parent box IDs for :attr:`partition_box_ids`. |
| 74 | +
|
| 75 | + .. attribute:: partition_parent_map |
| 76 | +
|
| 77 | + An object :class:`~numpy.ndarray`, where each entry maps a parent |
| 78 | + to its children indices in :attr:`partition_box_ids`. This can be used to |
| 79 | + :func:`cluster` all indices in ``partition_parent_map[i]`` into their |
| 80 | + parent. |
| 81 | +
|
| 82 | + .. automethod:: parent |
| 83 | + """ |
| 84 | + |
| 85 | + # level info |
| 86 | + level: int |
| 87 | + partition_box_ids: np.ndarray |
| 88 | + |
| 89 | + # tree info |
| 90 | + nlevels: int |
| 91 | + box_parent_ids: np.ndarray |
| 92 | + box_levels: np.ndarray |
| 93 | + |
| 94 | + _tree: Optional[Tree] |
| 95 | + |
| 96 | + @property |
| 97 | + def nclusters(self): |
| 98 | + return self.partition_box_ids.size |
| 99 | + |
| 100 | + @property |
| 101 | + def partition_parent_ids(self): |
| 102 | + return self.box_parent_ids[self.partition_box_ids] |
| 103 | + |
| 104 | + @property |
| 105 | + @memoize_method |
| 106 | + def partition_parent_map(self): |
| 107 | + # NOTE: np.unique returns a sorted array |
| 108 | + unique_parent_ids = np.unique(self.partition_parent_ids) |
| 109 | + # find the index of each parent id |
| 110 | + unique_parent_index = np.searchsorted( |
| 111 | + unique_parent_ids, self.partition_parent_ids |
| 112 | + ) |
| 113 | + |
| 114 | + unique_parent_map = np.empty(unique_parent_ids.size, dtype=object) |
| 115 | + for i in range(unique_parent_ids.size): |
| 116 | + unique_parent_map[i], = np.nonzero(unique_parent_index == i) |
| 117 | + |
| 118 | + return unique_parent_map |
| 119 | + |
| 120 | + def parent(self) -> "ClusterTreeLevel": |
| 121 | + """ |
| 122 | + :returns: a new :class:`ClusterTreeLevel` that represents the parent of |
| 123 | + the current one, with appropriately updated :attr:`partition_box_ids`, |
| 124 | + etc. |
| 125 | + """ |
| 126 | + |
| 127 | + if self.nclusters == 1: |
| 128 | + return self |
| 129 | + |
| 130 | + return replace(self, |
| 131 | + level=self.level - 1, |
| 132 | + partition_box_ids=np.unique(self.partition_parent_ids)) |
| 133 | + |
| 134 | + |
| 135 | +@singledispatch |
| 136 | +def cluster(obj: T, ctree: ClusterTreeLevel) -> T: |
| 137 | + """Merge together elements of *obj* into their parent object, as described |
| 138 | + by *ctree*. |
| 139 | + """ |
| 140 | + raise NotImplementedError(type(obj).__name__) |
| 141 | + |
| 142 | + |
| 143 | +@cluster.register(IndexList) |
| 144 | +def _cluster_index_list(obj: IndexList, ctree: ClusterTreeLevel) -> IndexList: |
| 145 | + assert obj.nclusters == ctree.nclusters |
| 146 | + |
| 147 | + if ctree.nclusters == 1: |
| 148 | + return obj |
| 149 | + |
| 150 | + sizes = [ |
| 151 | + sum(obj.cluster_size(j) for j in ppm) |
| 152 | + for ppm in ctree.partition_parent_map |
| 153 | + ] |
| 154 | + return replace(obj, starts=np.cumsum([0] + sizes)) |
| 155 | + |
| 156 | + |
| 157 | +@cluster.register(TargetAndSourceClusterList) |
| 158 | +def _cluster_target_and_source_cluster_list( |
| 159 | + obj: TargetAndSourceClusterList, ctree: ClusterTreeLevel, |
| 160 | + ) -> TargetAndSourceClusterList: |
| 161 | + assert obj.nclusters == ctree.nclusters |
| 162 | + |
| 163 | + if ctree.nclusters == 1: |
| 164 | + return obj |
| 165 | + |
| 166 | + return replace(obj, |
| 167 | + targets=cluster(obj.targets, ctree), |
| 168 | + sources=cluster(obj.sources, ctree)) |
| 169 | + |
| 170 | +# }}} |
| 171 | + |
| 172 | + |
| 173 | +# {{{ cluster generation |
| 174 | + |
| 175 | +def _build_binary_tree_from_indices(starts: np.ndarray) -> ClusterTreeLevel: |
| 176 | + return None |
| 177 | + |
| 178 | + |
| 179 | +def partition_by_nodes( |
| 180 | + actx: PyOpenCLArrayContext, places: GeometryCollection, *, |
| 181 | + dofdesc: Optional[sym.DOFDescriptorLike] = None, |
| 182 | + tree_kind: Optional[str] = "adaptive-level-restricted", |
| 183 | + max_particles_in_box: Optional[int] = None) -> IndexList: |
| 184 | + """Generate equally sized ranges of nodes. The partition is created at the |
| 185 | + lowest level of granularity, i.e. nodes. This results in balanced ranges |
| 186 | + of points, but will split elements across different ranges. |
| 187 | +
|
| 188 | + :arg dofdesc: a :class:`~pytential.symbolic.dof_desc.DOFDescriptor` for |
| 189 | + the geometry in *places* which should be partitioned. |
| 190 | + :arg tree_kind: if not *None*, it is passed to :class:`boxtree.TreeBuilder`. |
| 191 | + :arg max_particles_in_box: value used to control the number of points |
| 192 | + in each partition (and thus the number of partitions). See the documentation |
| 193 | + in :class:`boxtree.TreeBuilder`. |
| 194 | + """ |
| 195 | + if dofdesc is None: |
| 196 | + dofdesc = places.auto_source |
| 197 | + dofdesc = sym.as_dofdesc(dofdesc) |
| 198 | + |
| 199 | + if max_particles_in_box is None: |
| 200 | + max_particles_in_box = _DEFAULT_MAX_PARTICLES_IN_BOX |
| 201 | + |
| 202 | + lpot_source = places.get_geometry(dofdesc.geometry) |
| 203 | + discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage) |
| 204 | + |
| 205 | + if tree_kind is not None: |
| 206 | + from pytential.qbx.utils import tree_code_container |
| 207 | + tcc = tree_code_container(lpot_source._setup_actx) |
| 208 | + |
| 209 | + from arraycontext import flatten |
| 210 | + from meshmode.dof_array import DOFArray |
| 211 | + tree, _ = tcc.build_tree()(actx.queue, |
| 212 | + particles=flatten( |
| 213 | + actx.thaw(discr.nodes()), actx, leaf_class=DOFArray |
| 214 | + ), |
| 215 | + max_particles_in_box=max_particles_in_box, |
| 216 | + kind=tree_kind) |
| 217 | + |
| 218 | + from boxtree import box_flags_enum |
| 219 | + tree = tree.get(actx.queue) |
| 220 | + leaf_boxes, = (tree.box_flags & box_flags_enum.HAS_CHILDREN == 0).nonzero() |
| 221 | + |
| 222 | + indices = np.empty(len(leaf_boxes), dtype=object) |
| 223 | + starts = None |
| 224 | + |
| 225 | + for i, ibox in enumerate(leaf_boxes): |
| 226 | + box_start = tree.box_source_starts[ibox] |
| 227 | + box_end = box_start + tree.box_source_counts_cumul[ibox] |
| 228 | + indices[i] = tree.user_source_ids[box_start:box_end] |
| 229 | + |
| 230 | + ctree = ClusterTreeLevel( |
| 231 | + level=tree.nlevels - 1, |
| 232 | + nlevels=tree.nlevels, |
| 233 | + box_parent_ids=tree.box_parent_ids, |
| 234 | + box_levels=tree.box_levels, |
| 235 | + partition_box_ids=leaf_boxes, |
| 236 | + _tree=tree) |
| 237 | + else: |
| 238 | + if discr.ambient_dim != 2 and discr.dim == 1: |
| 239 | + raise ValueError("only curves are supported for 'tree_kind=None'") |
| 240 | + |
| 241 | + nclusters = max(discr.ndofs // max_particles_in_box, 2) |
| 242 | + indices = np.arange(0, discr.ndofs, dtype=np.int64) |
| 243 | + starts = np.linspace(0, discr.ndofs, nclusters + 1, dtype=np.int64) |
| 244 | + assert starts[-1] == discr.ndofs |
| 245 | + |
| 246 | + ctree = None |
| 247 | + |
| 248 | + from pytential.linalg import make_index_list |
| 249 | + return make_index_list(indices, starts=starts), ctree |
| 250 | + |
| 251 | +# }}} |
0 commit comments