add recursive clustering and skeletonization

Author: alexfikl

doc/conf.py

@@ -22,6 +22,10 @@
         "DOFDescriptorLike": "pytential.symbolic.dof_desc.DOFDescriptorLike",
         }
 
+nitpick_ignore_regex = [
+    ["py:class", r"_ProxyNeighborEvaluationResult"],
+    ]
+
 intersphinx_mapping = {
     "https://docs.python.org/3/": None,
     "https://numpy.org/doc/stable/": None,

doc/linalg.rst

@@ -8,7 +8,7 @@ scheme is used:
   component of the Stokeslet.
 * ``cluster`` refers to a piece of a ``block`` as used by the recursive
   proxy-based skeletonization of the direct solver algorithms. Clusters
-  are represented by a :class:`~pytential.linalg.TargetAndSourceClusterList`.
+  are represented by a :class:`~pytential.linalg.utils.TargetAndSourceClusterList`.
 
 GMRES
 -----
@@ -20,17 +20,31 @@ GMRES
 Hierarchical Direct Solver
 --------------------------
 
+.. note::
+
+    High-level API for direct solvers is in progress.
+
+Low-level Functionality
+-----------------------
+
 .. warning::
 
     All the classes and routines in this module are experimental and the
     API can change at any point.
 
+.. automodule:: pytential.linalg.skeletonization
+.. automodule:: pytential.linalg.cluster
 .. automodule:: pytential.linalg.proxy
-.. automodule:: pytential.linalg.utils
 
-Internal Functionality
-----------------------
+Internal Functionality and Utilities
+------------------------------------
+
+.. warning::
 
+    All the classes and routines in this module are experimental and the
+    API can change at any point.
+
+.. automodule:: pytential.linalg.utils
 .. automodule:: pytential.linalg.direct_solver_symbolic
 
 .. vim: sw=4:tw=75:fdm=marker

pytential/linalg/__init__.py

@@ -25,25 +25,9 @@
         make_index_list, make_index_cluster_cartesian_product,
         interp_decomp,
         )
-from pytential.linalg.proxy import (
-        ProxyClusterGeometryData, ProxyPointTarget, ProxyPointSource,
-        ProxyGeneratorBase, ProxyGenerator, QBXProxyGenerator,
-        partition_by_nodes, gather_cluster_neighbor_points,
-        )
-from pytential.linalg.skeletonization import (
-        SkeletonizationWrangler, make_skeletonization_wrangler,
-        SkeletonizationResult, skeletonize_by_proxy,
-        )
 
 __all__ = (
     "IndexList", "TargetAndSourceClusterList",
     "make_index_list", "make_index_cluster_cartesian_product",
     "interp_decomp",
-
-    "ProxyClusterGeometryData", "ProxyPointTarget", "ProxyPointSource",
-    "ProxyGeneratorBase", "ProxyGenerator", "QBXProxyGenerator",
-    "partition_by_nodes", "gather_cluster_neighbor_points",
-
-    "SkeletonizationWrangler", "make_skeletonization_wrangler",
-    "SkeletonizationResult", "skeletonize_by_proxy",
 )

pytential/linalg/cluster.py

@@ -0,0 +1,315 @@
+__copyright__ = "Copyright (C) 2022 Alexandru Fikl"
+
+__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.
+"""
+
+from dataclasses import dataclass, replace
+from functools import singledispatch
+from typing import Iterator, Optional
+
+import numpy as np
+
+from pytools import T
+
+from arraycontext import PyOpenCLArrayContext
+from boxtree.tree import Tree
+from pytential import sym, GeometryCollection
+from pytential.linalg.utils import IndexList, TargetAndSourceClusterList
+
+__doc__ = """
+Clustering
+~~~~~~~~~~
+
+.. autoclass:: ClusterLevel
+.. autoclass:: ClusterTree
+
+.. autofunction:: cluster
+.. autofunction:: partition_by_nodes
+"""
+
+# FIXME: this is just an arbitrary value
+_DEFAULT_MAX_PARTICLES_IN_BOX = 32
+
+
+# {{{ cluster tree
+
+
+def make_cluster_parent_map(parent_ids: np.ndarray) -> np.ndarray:
+    """Construct a parent map for :attr:`ClusterLevel.parent_map`."""
+    # NOTE: np.unique returns a sorted array
+    unique_parent_ids = np.unique(parent_ids)
+    # find the index of each parent id
+    unique_parent_index = np.searchsorted(unique_parent_ids, parent_ids)
+
+    unique_parent_map = np.empty(unique_parent_ids.size, dtype=object)
+    for i in range(unique_parent_ids.size):
+        unique_parent_map[i], = np.nonzero(unique_parent_index == i)
+
+    return unique_parent_map
+
+
+@dataclass(frozen=True)
+class ClusterLevel:
+    """A level in a :class:`ClusterTree`.
+
+    .. attribute:: level
+
+        Current level that is represented.
+
+    .. attribute:: nclusters
+
+        Number of clusters on the current level (same as number of boxes
+        in :attr:`box_ids`).
+
+    .. attribute:: box_ids
+
+        Box IDs on the current level.
+
+    .. attribute:: parent_map
+
+        An object :class:`~numpy.ndarray`, where each entry maps a parent
+        to its children indices in :attr:`box_ids`. This can be used to
+        :func:`cluster` all indices in ``parent_map[i]`` into their
+        parent.
+    """
+
+    level: int
+    box_ids: np.ndarray
+    parent_map: np.ndarray
+
+    @property
+    def nclusters(self):
+        return self.box_ids.size
+
+
+@dataclass(frozen=True)
+class ClusterTree:
+    """Hierarhical cluster representation.
+
+    .. attribute:: nlevels
+
+        Total number of levels in the tree.
+
+    .. attribute:: leaf_cluster_box_ids
+
+        Box IDs for each cluster on the leaf level of the tree.
+
+    .. attribute:: tree_cluster_parent_ids
+
+        Parent box IDs for :attr:`leaf_cluster_box_ids`.
+
+    .. automethod:: levels
+    """
+
+    nlevels: int
+    leaf_cluster_box_ids: np.ndarray
+    tree_cluster_parent_ids: np.ndarray
+
+    # NOTE: only here to allow easier debugging + testing
+    _tree: Optional[Tree]
+
+    @property
+    def nclusters(self):
+        return self.leaf_cluster_box_ids.size
+
+    def levels(self) -> Iterator[ClusterLevel]:
+        """
+        :returns: an iterator over all the :class:`ClusterLevel` levels.
+        """
+
+        box_ids = self.leaf_cluster_box_ids
+        parent_ids = self.tree_cluster_parent_ids[box_ids]
+        clevel = ClusterLevel(
+            level=self.nlevels - 1,
+            box_ids=box_ids,
+            parent_map=make_cluster_parent_map(parent_ids),
+            )
+
+        for _ in range(self.nlevels - 1, 0, -1):
+            yield clevel
+
+            box_ids = np.unique(self.tree_cluster_parent_ids[clevel.box_ids])
+            parent_ids = self.tree_cluster_parent_ids[box_ids]
+            clevel = ClusterLevel(
+                level=clevel.level - 1,
+                box_ids=box_ids,
+                parent_map=make_cluster_parent_map(parent_ids)
+                )
+
+        assert clevel.nclusters == 1
+
+
+@singledispatch
+def cluster(obj: T, ctree: ClusterTree) -> T:
+    """Merge together elements of *obj* into their parent object, as described
+    by *ctree*.
+    """
+    raise NotImplementedError(type(obj).__name__)
+
+
+@cluster.register(IndexList)
+def _cluster_index_list(obj: IndexList, clevel: ClusterLevel) -> IndexList:
+    assert obj.nclusters == clevel.nclusters
+
+    if clevel.nclusters == 1:
+        return obj
+
+    sizes = [sum([obj.cluster_size(j) for j in ppm]) for ppm in clevel.parent_map]
+    return replace(obj, starts=np.cumsum([0] + sizes))
+
+
+@cluster.register(TargetAndSourceClusterList)
+def _cluster_target_and_source_cluster_list(
+        obj: TargetAndSourceClusterList, clevel: ClusterLevel,
+        ) -> TargetAndSourceClusterList:
+    assert obj.nclusters == clevel.nclusters
+
+    if clevel.nclusters == 1:
+        return obj
+
+    return replace(obj,
+        targets=cluster(obj.targets, clevel),
+        sources=cluster(obj.sources, clevel))
+
+
+@cluster.register(np.ndarray)
+def _cluster_ndarray(obj: np.ndarray, clevel: ClusterLevel) -> np.ndarray:
+    assert obj.shape == (clevel.nclusters,)
+    assert all(block.ndim == 2 for block in obj)
+
+    if clevel.nclusters == 1:
+        return obj
+
+    def make_block(i: int, j: int):
+        if i == j:
+            return obj[i]
+
+        return np.zeros((obj[i].shape[0], obj[j].shape[1]), dtype=obj[i].dtype)
+
+    from pytools.obj_array import make_obj_array
+    return make_obj_array([
+        np.block([[make_block(i, j) for j in ppm] for i in ppm])
+        for ppm in clevel.parent_map
+        ])
+
+# }}}
+
+
+# {{{ cluster generation
+
+def _build_binary_ish_tree_from_indices(starts: np.ndarray) -> ClusterTree:
+    partition_box_ids = np.arange(starts.size - 1)
+
+    box_ids = partition_box_ids
+
+    box_parent_ids = []
+    offset = box_ids.size
+    while box_ids.size > 1:
+        # NOTE: this is probably not the most efficient way to do it, but this
+        # code is mostly meant for debugging using a simple tree
+        clusters = np.array_split(box_ids, box_ids.size // 2)
+        parent_ids = offset + np.arange(len(clusters))
+        box_parent_ids.append(np.repeat(parent_ids, [len(c) for c in clusters]))
+
+        box_ids = parent_ids
+        offset += box_ids.size
+
+    # NOTE: make the root point to itself
+    box_parent_ids.append(np.array([offset - 1]))
+    nlevels = len(box_parent_ids)
+
+    return ClusterTree(
+        nlevels=nlevels,
+        leaf_cluster_box_ids=partition_box_ids,
+        tree_cluster_parent_ids=np.concatenate(box_parent_ids),
+        _tree=None)
+
+
+def partition_by_nodes(
+        actx: PyOpenCLArrayContext, places: GeometryCollection, *,
+        dofdesc: Optional[sym.DOFDescriptorLike] = None,
+        tree_kind: Optional[str] = "adaptive-level-restricted",
+        max_particles_in_box: Optional[int] = None) -> IndexList:
+    """Generate equally sized ranges of nodes. The partition is created at the
+    lowest level of granularity, i.e. nodes. This results in balanced ranges
+    of points, but will split elements across different ranges.
+
+    :arg dofdesc: a :class:`~pytential.symbolic.dof_desc.DOFDescriptor` for
+        the geometry in *places* which should be partitioned.
+    :arg tree_kind: if not *None*, it is passed to :class:`boxtree.TreeBuilder`.
+    :arg max_particles_in_box: value used to control the number of points
+        in each partition (and thus the number of partitions). See the documentation
+        in :class:`boxtree.TreeBuilder`.
+    """
+    if dofdesc is None:
+        dofdesc = places.auto_source
+    dofdesc = sym.as_dofdesc(dofdesc)
+
+    if max_particles_in_box is None:
+        max_particles_in_box = _DEFAULT_MAX_PARTICLES_IN_BOX
+
+    lpot_source = places.get_geometry(dofdesc.geometry)
+    discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage)
+
+    if tree_kind is not None:
+        from pytential.qbx.utils import tree_code_container
+        tcc = tree_code_container(lpot_source._setup_actx)
+
+        from arraycontext import flatten
+        from meshmode.dof_array import DOFArray
+        tree, _ = tcc.build_tree()(actx.queue,
+                particles=flatten(
+                    actx.thaw(discr.nodes()), actx, leaf_class=DOFArray
+                    ),
+                max_particles_in_box=max_particles_in_box,
+                kind=tree_kind)
+
+        from boxtree import box_flags_enum
+        tree = tree.get(actx.queue)
+        leaf_boxes, = (tree.box_flags & box_flags_enum.HAS_CHILDREN == 0).nonzero()
+
+        indices = np.empty(len(leaf_boxes), dtype=object)
+        starts = None
+
+        for i, ibox in enumerate(leaf_boxes):
+            box_start = tree.box_source_starts[ibox]
+            box_end = box_start + tree.box_source_counts_cumul[ibox]
+            indices[i] = tree.user_source_ids[box_start:box_end]
+
+        ctree = ClusterTree(
+            nlevels=tree.nlevels,
+            leaf_cluster_box_ids=leaf_boxes,
+            tree_cluster_parent_ids=tree.box_parent_ids,
+            _tree=tree)
+    else:
+        if discr.ambient_dim != 2 and discr.dim == 1:
+            raise ValueError("only curves are supported for 'tree_kind=None'")
+
+        nclusters = max(discr.ndofs // max_particles_in_box, 2)
+        indices = np.arange(0, discr.ndofs, dtype=np.int64)
+        starts = np.linspace(0, discr.ndofs, nclusters + 1, dtype=np.int64)
+        assert starts[-1] == discr.ndofs
+
+        ctree = _build_binary_ish_tree_from_indices(starts)
+
+    from pytential.linalg import make_index_list
+    return make_index_list(indices, starts=starts), ctree
+
+# }}}