diff --git a/doc/conf.py b/doc/conf.py index 0eb9e4d82..5b76c6199 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -22,6 +22,10 @@ "DOFDescriptorLike": "pytential.symbolic.dof_desc.DOFDescriptorLike", } +nitpick_ignore_regex = [ + ["py:class", r"_ProxyNeighborEvaluationResult"], + ] + intersphinx_mapping = { "arraycontext": ("https://documen.tician.de/arraycontext", None), "boxtree": ("https://documen.tician.de/boxtree", None), diff --git a/doc/linalg.rst b/doc/linalg.rst index 97497434f..d3d4870ce 100644 --- a/doc/linalg.rst +++ b/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 diff --git a/pytential/linalg/__init__.py b/pytential/linalg/__init__.py index bfbd4d82e..eed539cc8 100644 --- a/pytential/linalg/__init__.py +++ b/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", ) diff --git a/pytential/linalg/cluster.py b/pytential/linalg/cluster.py new file mode 100644 index 000000000..ca7291307 --- /dev/null +++ b/pytential/linalg/cluster.py @@ -0,0 +1,397 @@ +__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 Any, Iterator, Optional, Tuple + +import numpy as np + +from pytools import memoize_method +from pytools.obj_array import make_obj_array + +from arraycontext import PyOpenCLArrayContext +from boxtree.tree import Tree +from meshmode.discretization import Discretization +from pytential import sym, GeometryCollection +from pytential.linalg.utils import IndexList, TargetAndSourceClusterList +from pytential.qbx import QBXLayerPotentialSource + +__doc__ = """ +Clustering +~~~~~~~~~~ + +.. autoclass:: ClusterLevel +.. autoclass:: ClusterTree + +.. autofunction:: split_array +.. autofunction:: cluster +.. autofunction:: uncluster + +.. 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) + ids = np.arange(parent_ids.size) + + return make_obj_array([ + ids[parent_ids == unique_parent_ids[i]] + for i in range(unique_parent_ids.size) + ]) + + +@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: + r"""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`. + + .. attribute:: levels + + An :class:`~numpy.ndarray` of :class:`ClusterLevel`\ s. + """ + + 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 + + @property + @memoize_method + def levels(self) -> np.ndarray: + return make_obj_array(list(self.iter_levels())) + + def iter_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, -1, -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 + +# }}} + + +# {{{ cluster + +def split_array(x: np.ndarray, index: IndexList) -> np.ndarray: + """ + :returns: an object :class:`~numpy.ndarray` where each entry contains the + elements of the :math:`i`-th cluster in *index*. + """ + assert x.size == index.nindices + + from pytools.obj_array import make_obj_array + return make_obj_array([ + index.cluster_take(x, i) for i in range(index.nclusters) + ]) + + +@singledispatch +def cluster(obj: Any, clevel: ClusterLevel) -> Any: + """Merge together elements of *obj* into their parent object, as described + by :attr:`ClusterLevel.parent_map`. + """ + 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 + + from pytential.linalg.utils import make_index_list + indices = make_obj_array([ + np.concatenate([obj.cluster_indices(i) for i in ppm]) + for ppm in clevel.parent_map + ]) + + return make_index_list(indices) + + +@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,) + 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 import single_valued + ndim = single_valued(block.ndim for block in obj) + + if ndim == 1: + return make_obj_array([ + np.concatenate([obj[i] for i in ppm]) for ppm in clevel.parent_map + ]) + elif ndim == 2: + return make_obj_array([ + np.block([[make_block(i, j) for j in ppm] for i in ppm]) + for ppm in clevel.parent_map + ]) + else: + raise ValueError(f"unsupported ndarray dimension: '{ndim}'") + +# }}} + + +# {{{ uncluster + +def uncluster(ary: np.ndarray, index: IndexList, clevel: ClusterLevel) -> np.ndarray: + """Performs the reverse of :func:`cluster` on object arrays. + + :arg ary: an object :class:`~numpy.ndarray` with a shape that matches + :attr:`ClusterLevel.parent_map`. + :arg index: an :class:`~pytential.linalg.utils.IndexList` for the + current level, as given by :attr:`ClusterLevel.box_ids`. + :returns: an object :class:`~numpy.ndarray` with a shape that matches + :attr:`ClusterLevel.box_ids` of all the elements of *ary* that belong + to each child cluster. + """ + assert ary.dtype.char == "O" + assert ary.shape == (clevel.parent_map.size,) + + if index.nclusters == 1: + return ary + + result: np.ndarray = np.empty(index.nclusters, dtype=object) + for ifrom, ppm in enumerate(clevel.parent_map): + offset = 0 + for ito in ppm: + cluster_size = index.cluster_size(ito) + result[ito] = ary[ifrom][offset:offset + cluster_size] + offset += cluster_size + + assert ary[ifrom].shape == (offset,) + + return result + +# }}} + + +# {{{ cluster generation + +def _build_binary_ish_tree_from_starts(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) -> Tuple[IndexList, ClusterTree]: + """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) + assert isinstance(lpot_source, QBXLayerPotentialSource) + + discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage) + assert isinstance(discr, Discretization) + + 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) + # FIXME maybe this should use IS_LEAF once available? + leaf_boxes, = ( + tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0 + ).nonzero() + + indices: np.ndarray = np.empty(len(leaf_boxes), dtype=object) + starts: Optional[np.ndarray] = 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_starts(starts) + + from pytential.linalg import make_index_list + return make_index_list(indices, starts=starts), ctree + +# }}} + +# vim: foldmethod=marker diff --git a/pytential/linalg/direct_solver_symbolic.py b/pytential/linalg/direct_solver_symbolic.py index 0215e694a..c57752840 100644 --- a/pytential/linalg/direct_solver_symbolic.py +++ b/pytential/linalg/direct_solver_symbolic.py @@ -20,6 +20,8 @@ THE SOFTWARE. """ +import numpy as np + from pytools.obj_array import make_obj_array from pytential.symbolic.mappers import ( @@ -44,9 +46,12 @@ class PROXY_SKELETONIZATION_TARGET: # noqa: N801 def prepare_expr(places, exprs, auto_where=None): from pytential.symbolic.execution import _prepare_expr - return make_obj_array([ - _prepare_expr(places, expr, auto_where=auto_where) - for expr in exprs]) + if isinstance(exprs, np.ndarray): + return make_obj_array([ + _prepare_expr(places, expr, auto_where=auto_where) + for expr in exprs]) + + return _prepare_expr(places, exprs, auto_where=auto_where) def prepare_proxy_expr(places, exprs, auto_where=None): @@ -222,3 +227,5 @@ def __init__(self, source, target): self.operand_rec = _LocationReplacer(source, default_source=source) # }}} + +# vim: foldmethod=marker diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py index bca1cfb88..0f523cf91 100644 --- a/pytential/linalg/proxy.py +++ b/pytential/linalg/proxy.py @@ -46,8 +46,6 @@ Proxy Point Generation ~~~~~~~~~~~~~~~~~~~~~~ -.. currentmodule:: pytential.linalg - .. autoclass:: ProxyPointSource .. autoclass:: ProxyPointTarget .. autoclass:: ProxyClusterGeometryData @@ -56,7 +54,6 @@ .. autoclass:: ProxyGenerator .. autoclass:: QBXProxyGenerator -.. autofunction:: partition_by_nodes .. autofunction:: gather_cluster_neighbor_points """ @@ -64,79 +61,6 @@ _DEFAULT_MAX_PARTICLES_IN_BOX = 32 -# {{{ point index partitioning - -def partition_by_nodes( - actx: PyOpenCLArrayContext, places: GeometryCollection, *, - dofdesc: Optional["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 - - from pytential.source import LayerPotentialSourceBase - - lpot_source = places.get_geometry(dofdesc.geometry) - assert isinstance(lpot_source, LayerPotentialSourceBase) - - discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage) - assert isinstance(discr, Discretization) - - if tree_kind is not None: - from pytential.qbx.utils import tree_code_container - tcc = tree_code_container(lpot_source._setup_actx) - - 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) - # FIXME maybe this should use IS_LEAF once available? - leaf_boxes, = ( - tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0 - ).nonzero() - - indices: np.ndarray = np.empty(len(leaf_boxes), dtype=object) - starts: Optional[np.ndarray] = 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] - 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 - - from pytential.linalg import make_index_list - return make_index_list(indices, starts=starts) - -# }}} - - # {{{ proxy points class ProxyPointSource(PointPotentialSource): @@ -187,12 +111,12 @@ class ProxyClusterGeometryData: """ .. attribute:: srcindex - A :class:`~pytential.linalg.IndexList` describing which cluster + A :class:`~pytential.linalg.utils.IndexList` describing which cluster of points each proxy ball was created from. .. attribute:: pxyindex - A :class:`~pytential.linalg.IndexList` describing which proxies + A :class:`~pytential.linalg.utils.IndexList` describing which proxies belong to which cluster. .. attribute:: points diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py index b36dbba0e..c3076610a 100644 --- a/pytential/linalg/skeletonization.py +++ b/pytential/linalg/skeletonization.py @@ -22,31 +22,34 @@ from dataclasses import dataclass from typing import ( - Any, Callable, Dict, Hashable, Optional, Sequence, Tuple, Union) + Any, Callable, Dict, Hashable, Iterable, Optional, Sequence, Tuple, Union, + TYPE_CHECKING) import numpy as np from arraycontext import PyOpenCLArrayContext, Array from pytential import GeometryCollection, sym -from pytential.linalg.utils import IndexList, TargetAndSourceClusterList from pytential.linalg.proxy import ProxyGeneratorBase, ProxyClusterGeometryData +from pytential.linalg.cluster import ClusterTree, cluster from pytential.linalg.direct_solver_symbolic import ( PROXY_SKELETONIZATION_TARGET, PROXY_SKELETONIZATION_SOURCE, prepare_expr, prepare_proxy_expr) +if TYPE_CHECKING: + from pytential.linalg.utils import IndexList, TargetAndSourceClusterList -__doc__ = """ -.. currentmodule:: pytential.linalg +__doc__ = """ Skeletonization ---------------- +~~~~~~~~~~~~~~~ .. autoclass:: SkeletonizationWrangler .. autoclass:: make_skeletonization_wrangler .. autoclass:: SkeletonizationResult .. autofunction:: skeletonize_by_proxy +.. autofunction:: rec_skeletonize_by_proxy """ @@ -156,8 +159,8 @@ def _apply_weights( actx: PyOpenCLArrayContext, mat: np.ndarray, places: GeometryCollection, - tgt_pxy_index: TargetAndSourceClusterList, - cluster_index: IndexList, + tgt_pxy_index: "TargetAndSourceClusterList", + cluster_index: "IndexList", domain: sym.DOFDescriptor) -> np.ndarray: """Computes the weights using :func:`_approximate_geometry_waa_magnitude` and multiplies each cluster in *mat* by its corresponding weight. @@ -233,13 +236,13 @@ class SkeletonizationWrangler: A callable that is used to evaluate farfield proxy interactions. This should follow the calling convention of the constructor to - :class:`pytential.symbolic.matrix.P2PClusterMatrixBuilder`. + ``pytential.symbolic.matrix.P2PClusterMatrixBuilder``. .. attribute:: neighbor_cluster_builder A callable that is used to evaluate nearfield neighbour interactions. This should follow the calling convention of the constructor to - :class:`pytential.symbolic.matrix.QBXClusterMatrixBuilder`. + ``pytential.symbolic.matrix.QBXClusterMatrixBuilder``. .. automethod:: evaluate_source_proxy_interaction .. automethod:: evaluate_target_proxy_interaction @@ -290,13 +293,28 @@ def _evaluate_expr(self, context=self.context, **kwargs)(expr) + def evaluate_self( + self, + actx: PyOpenCLArrayContext, + places: GeometryCollection, + tgt_src_index: "TargetAndSourceClusterList", + ibrow: int, ibcol: int, + ) -> np.ndarray: + cls = self.neighbor_cluster_builder + return self._evaluate_expr( + actx, places, cls, tgt_src_index, self.exprs[ibrow], + idomain=ibcol, _weighted=self.weighted_sources) + # {{{ nearfield - def evaluate_source_neighbor_interaction(self, + def evaluate_source_neighbor_interaction( + self, actx: PyOpenCLArrayContext, places: GeometryCollection, - pxy: ProxyClusterGeometryData, nbrindex: IndexList, *, - ibrow: int, ibcol: int) -> Tuple[np.ndarray, TargetAndSourceClusterList]: + pxy: ProxyClusterGeometryData, nbrindex: "IndexList", *, + ibrow: int, ibcol: int, + ) -> Tuple[np.ndarray, "TargetAndSourceClusterList"]: + from pytential.linalg.utils import TargetAndSourceClusterList nbr_src_index = TargetAndSourceClusterList(nbrindex, pxy.srcindex) eval_mapper_cls = self.neighbor_cluster_builder @@ -307,11 +325,14 @@ def evaluate_source_neighbor_interaction(self, return mat, nbr_src_index - def evaluate_target_neighbor_interaction(self, + def evaluate_target_neighbor_interaction( + self, actx: PyOpenCLArrayContext, places: GeometryCollection, - pxy: ProxyClusterGeometryData, nbrindex: IndexList, *, - ibrow: int, ibcol: int) -> Tuple[np.ndarray, TargetAndSourceClusterList]: + pxy: ProxyClusterGeometryData, nbrindex: "IndexList", *, + ibrow: int, ibcol: int, + ) -> Tuple[np.ndarray, "TargetAndSourceClusterList"]: + from pytential.linalg.utils import TargetAndSourceClusterList tgt_nbr_index = TargetAndSourceClusterList(pxy.srcindex, nbrindex) eval_mapper_cls = self.neighbor_cluster_builder @@ -326,13 +347,17 @@ def evaluate_target_neighbor_interaction(self, # {{{ proxy - def evaluate_source_proxy_interaction(self, + def evaluate_source_proxy_interaction( + self, actx: PyOpenCLArrayContext, places: GeometryCollection, - pxy: ProxyClusterGeometryData, nbrindex: IndexList, *, - ibrow: int, ibcol: int) -> Tuple[np.ndarray, TargetAndSourceClusterList]: - from pytential.collection import add_geometry_to_collection + pxy: ProxyClusterGeometryData, nbrindex: "IndexList", *, + ibrow: int, ibcol: int, + ) -> Tuple[np.ndarray, "TargetAndSourceClusterList"]: + from pytential.linalg.utils import TargetAndSourceClusterList pxy_src_index = TargetAndSourceClusterList(pxy.pxyindex, pxy.srcindex) + + from pytential.collection import add_geometry_to_collection places = add_geometry_to_collection( places, {PROXY_SKELETONIZATION_TARGET: pxy.as_targets()} ) @@ -347,13 +372,17 @@ def evaluate_source_proxy_interaction(self, return mat, pxy_src_index - def evaluate_target_proxy_interaction(self, + def evaluate_target_proxy_interaction( + self, actx: PyOpenCLArrayContext, places: GeometryCollection, - pxy: ProxyClusterGeometryData, nbrindex: IndexList, *, - ibrow: int, ibcol: int) -> Tuple[np.ndarray, TargetAndSourceClusterList]: - from pytential.collection import add_geometry_to_collection + pxy: ProxyClusterGeometryData, nbrindex: "IndexList", *, + ibrow: int, ibcol: int, + ) -> Tuple[np.ndarray, "TargetAndSourceClusterList"]: + from pytential.linalg.utils import TargetAndSourceClusterList tgt_pxy_index = TargetAndSourceClusterList(pxy.srcindex, pxy.pxyindex) + + from pytential.collection import add_geometry_to_collection places = add_geometry_to_collection( places, {PROXY_SKELETONIZATION_SOURCE: pxy.as_sources()} ) @@ -510,10 +539,10 @@ class _ProxyNeighborEvaluationResult: pxy: ProxyClusterGeometryData pxymat: np.ndarray - pxyindex: TargetAndSourceClusterList + pxyindex: "TargetAndSourceClusterList" nbrmat: np.ndarray - nbrindex: TargetAndSourceClusterList + nbrindex: "TargetAndSourceClusterList" def __getitem__(self, i: int) -> Tuple[np.ndarray, np.ndarray]: """ @@ -536,11 +565,11 @@ def _evaluate_proxy_skeletonization_interaction( places: GeometryCollection, proxy_generator: ProxyGeneratorBase, wrangler: SkeletonizationWrangler, - cluster_index: IndexList, *, + cluster_index: "IndexList", *, evaluate_proxy: Callable[..., - Tuple[np.ndarray, TargetAndSourceClusterList]], + Tuple[np.ndarray, "TargetAndSourceClusterList"]], evaluate_neighbor: Callable[..., - Tuple[np.ndarray, TargetAndSourceClusterList]], + Tuple[np.ndarray, "TargetAndSourceClusterList"]], dofdesc: Optional[sym.DOFDescriptor] = None, max_particles_in_box: Optional[int] = None, ) -> _ProxyNeighborEvaluationResult: @@ -551,7 +580,7 @@ def _evaluate_proxy_skeletonization_interaction( if cluster_index.nclusters == 1: raise ValueError("cannot make a proxy skeleton for a single cluster") - from pytential.linalg import gather_cluster_neighbor_points + from pytential.linalg.proxy import gather_cluster_neighbor_points pxy = proxy_generator(actx, dofdesc, cluster_index) nbrindex = gather_cluster_neighbor_points( actx, pxy, @@ -571,7 +600,7 @@ def _skeletonize_block_by_proxy_with_mats( places: GeometryCollection, proxy_generator: ProxyGeneratorBase, wrangler: SkeletonizationWrangler, - tgt_src_index: TargetAndSourceClusterList, *, + tgt_src_index: "TargetAndSourceClusterList", *, id_eps: Optional[float] = None, id_rank: Optional[int] = None, max_particles_in_box: Optional[int] = None ) -> "SkeletonizationResult": @@ -606,8 +635,8 @@ def _skeletonize_block_by_proxy_with_mats( ibrow=ibrow, ibcol=ibcol) ) - src_skl_indices: np.ndarray = np.empty(nclusters, dtype=object) - tgt_skl_indices: np.ndarray = np.empty(nclusters, dtype=object) + skel_src_indices: np.ndarray = np.empty(nclusters, dtype=object) + skel_tgt_indices: np.ndarray = np.empty(nclusters, dtype=object) skel_starts: np.ndarray = np.zeros(nclusters + 1, dtype=np.int32) L: np.ndarray = np.empty(nclusters, dtype=object) @@ -635,7 +664,7 @@ def _skeletonize_block_by_proxy_with_mats( interp = interp[:k, :] L[i] = interp.T - tgt_skl_indices[i] = tgt_src_index.targets.cluster_indices(i)[idx[:k]] + skel_tgt_indices[i] = tgt_src_index.targets.cluster_indices(i)[idx[:k]] assert L[i].shape == (tgt_mat.shape[0], k) # skeletonize source points @@ -643,22 +672,49 @@ def _skeletonize_block_by_proxy_with_mats( assert 0 < k <= len(idx) R[i] = interp - src_skl_indices[i] = tgt_src_index.sources.cluster_indices(i)[idx[:k]] + skel_src_indices[i] = tgt_src_index.sources.cluster_indices(i)[idx[:k]] assert R[i].shape == (k, src_mat.shape[1]) skel_starts[i + 1] = skel_starts[i] + k - assert tgt_skl_indices[i].shape == src_skl_indices[i].shape + assert skel_tgt_indices[i].shape == skel_src_indices[i].shape - from pytential.linalg import make_index_list - src_skl_index = make_index_list(np.hstack(list(src_skl_indices)), skel_starts) - tgt_skl_index = make_index_list(np.hstack(list(tgt_skl_indices)), skel_starts) - skel_tgt_src_index = TargetAndSourceClusterList(tgt_skl_index, src_skl_index) + # evaluate diagonal + from pytential.linalg.utils import make_flat_cluster_diag + mat = wrangler.evaluate_self(actx, places, tgt_src_index, ibrow, ibcol) + D = make_flat_cluster_diag(mat, tgt_src_index) + + from pytential.linalg import TargetAndSourceClusterList, make_index_list + skel_src_index = make_index_list(np.hstack(list(skel_src_indices)), skel_starts) + skel_tgt_index = make_index_list(np.hstack(list(skel_tgt_indices)), skel_starts) + skel_tgt_src_index = TargetAndSourceClusterList(skel_tgt_index, skel_src_index) return SkeletonizationResult( - L=L, R=R, + L=L, R=R, D=D, tgt_src_index=tgt_src_index, skel_tgt_src_index=skel_tgt_src_index, _src_eval_result=src_result, _tgt_eval_result=tgt_result) + +def _evaluate_root( + actx: PyOpenCLArrayContext, ibrow: int, ibcol: int, + places: GeometryCollection, + wrangler: SkeletonizationWrangler, + tgt_src_index: "TargetAndSourceClusterList" + ) -> "SkeletonizationResult": + assert tgt_src_index.nclusters == 1 + + from pytential.linalg.utils import make_flat_cluster_diag + mat = wrangler.evaluate_self(actx, places, tgt_src_index, ibrow, ibcol) + D = make_flat_cluster_diag(mat, tgt_src_index) + + from pytools.obj_array import make_obj_array + return SkeletonizationResult( + L=make_obj_array([np.eye(*D[0].shape)]), + R=make_obj_array([np.eye(*D[0].shape)]), + D=D, + tgt_src_index=tgt_src_index, skel_tgt_src_index=tgt_src_index, + _src_eval_result=None, _tgt_eval_result=None, + ) + # }}} @@ -697,29 +753,35 @@ class SkeletonizationResult: An object :class:`~numpy.ndarray` of size ``(nclusters,)``. + .. attribute:: D + + An object :class:`~numpy.ndarray` of size ``(nclusters,)``. This + contains the diagonal blocks for :attr:`tgt_src_index`. + .. attribute:: tgt_src_index - A :class:`~pytential.linalg.TargetAndSourceClusterList` representing the - indices in the original matrix :math:`A` that have been skeletonized. + A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` representing + the indices in the original matrix :math:`A` that has been skeletonized. .. attribute:: skel_tgt_src_index - A :class:`~pytential.linalg.TargetAndSourceClusterList` representing a - subset of :attr:`tgt_src_index`, i.e. the skeleton of each cluster of + A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` representing + a subset of :attr:`tgt_src_index`, i.e. the skeleton of each cluster of :math:`A`. These indices can be used to reconstruct the :math:`S` matrix. """ L: np.ndarray R: np.ndarray + D: np.ndarray - tgt_src_index: TargetAndSourceClusterList - skel_tgt_src_index: TargetAndSourceClusterList + tgt_src_index: "TargetAndSourceClusterList" + skel_tgt_src_index: "TargetAndSourceClusterList" # NOTE: these are meant only for testing! They contain the interactions # between the source / target points and their proxies / neighbors. - _src_eval_result: Optional[_ProxyNeighborEvaluationResult] = None - _tgt_eval_result: Optional[_ProxyNeighborEvaluationResult] = None + _src_eval_result: Optional["_ProxyNeighborEvaluationResult"] = None + _tgt_eval_result: Optional["_ProxyNeighborEvaluationResult"] = None def __post_init__(self): if __debug__: @@ -746,11 +808,12 @@ def skeletonize_by_proxy( actx: PyOpenCLArrayContext, places: GeometryCollection, - tgt_src_index: TargetAndSourceClusterList, - exprs: Union[sym.var, Sequence[sym.var]], - input_exprs: Union[sym.var, Sequence[sym.var]], *, + tgt_src_index: "TargetAndSourceClusterList", + exprs: Union[sym.var, Iterable[sym.var]], + input_exprs: Union[sym.var, Iterable[sym.var]], *, domains: Optional[Sequence[Hashable]] = None, context: Optional[Dict[str, Any]] = None, + auto_where: Optional[Any] = None, approx_nproxy: Optional[int] = None, proxy_radius_factor: Optional[float] = None, @@ -760,7 +823,7 @@ def skeletonize_by_proxy( max_particles_in_box: Optional[int] = None) -> np.ndarray: r"""Evaluate and skeletonize a symbolic expression using proxy-based methods. - :arg tgt_src_index: a :class:`~pytential.linalg.TargetAndSourceClusterList` + :arg tgt_src_index: a :class:`~pytential.linalg.utils.TargetAndSourceClusterList` indicating which indices participate in the skeletonization. :arg exprs: see :func:`make_skeletonization_wrangler`. @@ -768,8 +831,8 @@ def skeletonize_by_proxy( :arg domains: see :func:`make_skeletonization_wrangler`. :arg context: see :func:`make_skeletonization_wrangler`. - :arg approx_nproxy: see :class:`~pytential.linalg.ProxyGenerator`. - :arg proxy_radius_factor: see :class:`~pytential.linalg.ProxyGenerator`. + :arg approx_nproxy: see :class:`~pytential.linalg.proxy.ProxyGenerator`. + :arg proxy_radius_factor: see :class:`~pytential.linalg.proxy.ProxyGenerator`. :arg id_eps: a floating point value used as a tolerance in skeletonizing each block in *tgt_src_index*. @@ -783,19 +846,101 @@ def skeletonize_by_proxy( from pytential.linalg.proxy import QBXProxyGenerator wrangler = make_skeletonization_wrangler( places, exprs, input_exprs, - domains=domains, context=context) + domains=domains, context=context, auto_where=auto_where) proxy = QBXProxyGenerator(places, approx_nproxy=approx_nproxy, radius_factor=proxy_radius_factor) + if wrangler.nrows != 1 or wrangler.ncols != 1: + raise NotImplementedError("support for block matrices") + + from itertools import product skels: np.ndarray = np.empty((wrangler.nrows, wrangler.ncols), dtype=object) - for ibrow in range(wrangler.nrows): - for ibcol in range(wrangler.ncols): - skels[ibrow, ibcol] = _skeletonize_block_by_proxy_with_mats( - actx, ibrow, ibcol, places, proxy, wrangler, tgt_src_index, - id_eps=id_eps, id_rank=id_rank, - max_particles_in_box=max_particles_in_box) + for ibrow, ibcol in product(range(wrangler.nrows), range(wrangler.ncols)): + skels[ibrow, ibcol] = _skeletonize_block_by_proxy_with_mats( + actx, ibrow, ibcol, proxy.places, proxy, wrangler, tgt_src_index, + id_eps=id_eps, id_rank=id_rank, + max_particles_in_box=max_particles_in_box) return skels # }}} + + +# {{{ recursive skeletonization by proxy + +def rec_skeletonize_by_proxy( + actx: PyOpenCLArrayContext, + places: GeometryCollection, + + ctree: ClusterTree, + tgt_src_index: "TargetAndSourceClusterList", + exprs: Union[sym.var, Iterable[sym.var]], + input_exprs: Union[sym.var, Iterable[sym.var]], *, + domains: Optional[Sequence[Hashable]] = None, + context: Optional[Dict[str, Any]] = None, + auto_where: Optional[Any] = None, + + approx_nproxy: Optional[int] = None, + proxy_radius_factor: Optional[float] = None, + + id_eps: Optional[float] = None, + max_particles_in_box: Optional[int] = None, + + _wrangler: Optional[SkeletonizationWrangler] = None, + _proxy: Optional[ProxyGeneratorBase] = None) -> np.ndarray: + r"""Performs recursive skeletonization based on :func:`skeletonize_by_proxy`. + + :returns: an object :class:`~numpy.ndarray` of :class:`SkeletonizationResult`\ s, + one per level in *ctree*. + """ + + assert ctree.nclusters == tgt_src_index.nclusters + + if id_eps is None: + id_eps = 1.0e-8 + + if _proxy is None: + from pytential.linalg.proxy import QBXProxyGenerator + proxy: ProxyGeneratorBase = QBXProxyGenerator(places, + approx_nproxy=approx_nproxy, + radius_factor=proxy_radius_factor) + else: + proxy = _proxy + + if _wrangler is None: + wrangler = make_skeletonization_wrangler( + places, exprs, input_exprs, + domains=domains, context=context, auto_where=auto_where) + else: + wrangler = _wrangler + + if wrangler.nrows != 1 or wrangler.ncols != 1: + raise NotImplementedError("support for block matrices") + + from itertools import product + skel_per_level: np.ndarray = np.empty(ctree.nlevels, dtype=object) + for i, clevel in enumerate(ctree.levels[:-1]): + for ibrow, ibcol in product(range(wrangler.nrows), range(wrangler.ncols)): + skeleton = _skeletonize_block_by_proxy_with_mats( + actx, ibrow, ibcol, proxy.places, proxy, wrangler, tgt_src_index, + id_eps=id_eps, + # NOTE: we probably never want to set the rank here? + id_rank=None, + max_particles_in_box=max_particles_in_box) + + skel_per_level[i] = skeleton + tgt_src_index = cluster(skeleton.skel_tgt_src_index, clevel) + + assert tgt_src_index.nclusters == 1 + assert not isinstance(skel_per_level[-1], SkeletonizationResult) + + # evaluate the full root cluster (no skeletonization or anything) + skeleton = _evaluate_root(actx, 0, 0, places, wrangler, tgt_src_index) + skel_per_level[-1] = skeleton + + return skel_per_level + +# }}} + +# vim: foldmethod=marker diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py index beb671cd8..98bbe0ad1 100644 --- a/pytential/linalg/utils.py +++ b/pytential/linalg/utils.py @@ -34,16 +34,14 @@ __doc__ = """ -Misc -~~~~ - -.. currentmodule:: pytential.linalg - .. autoclass:: IndexList .. autoclass:: TargetAndSourceClusterList .. autofunction:: make_index_list .. autofunction:: make_index_cluster_cartesian_product +.. autofunction:: make_flat_cluster_diag + +.. autofunction:: interp_decomp """ @@ -54,6 +52,7 @@ class IndexList: """Convenience class for working with clusters (subsets) of an array. .. attribute:: nclusters + .. attribute:: nindices .. attribute:: indices An :class:`~numpy.ndarray` of not necessarily continuous or increasing @@ -78,6 +77,10 @@ class IndexList: def nclusters(self) -> int: return self.starts.size - 1 + @property + def nindices(self) -> int: + return self.indices.size + def cluster_size(self, i: int) -> int: if not (0 <= i < self.nclusters): raise IndexError( @@ -113,6 +116,11 @@ class TargetAndSourceClusterList: """Convenience class for working with clusters (subsets) of a matrix. .. attribute:: nclusters + + .. attribute:: shape + + Shape of the Cartesian product of the :attr:`targets` and :attr:`sources`. + .. attribute:: targets An :class:`IndexList` encapsulating target cluster indices. @@ -153,6 +161,13 @@ def _flat_cluster_starts(self): def _flat_total_size(self): return self._flat_cluster_starts[-1] + def __iter__(self): + return iter((self.targets, self.sources)) + + @property + def shape(self): + return (self.targets.nindices, self.sources.nindices) + def cluster_shape(self, i: int, j: int) -> Tuple[int, int]: r""" :returns: the shape of the cluster ``(i, j)``, where *i* indexes into @@ -308,15 +323,14 @@ def make_flat_cluster_diag( correspondence to the index sets constructed by :func:`make_index_cluster_cartesian_product` for *mindex*. - :returns: a block diagonal object :class:`~numpy.ndarray`, where each - diagonal element :math:`(i, i)` is the reshaped slice of *mat* that - corresponds to the cluster :math:`i`. + :returns: an object :class:`~numpy.ndarray`, where each element represents + the block of a block-diagonal matrix. """ - cluster_mat: np.ndarray = np.full((mindex.nclusters, mindex.nclusters), 0, - dtype=object) + + cluster_mat: np.ndarray = np.empty(mindex.nclusters, dtype=object) for i in range(mindex.nclusters): shape = mindex.cluster_shape(i, i) - cluster_mat[i, i] = mindex.flat_cluster_take(mat, i).reshape(*shape) + cluster_mat[i] = mindex.flat_cluster_take(mat, i).reshape(*shape) return cluster_mat @@ -334,8 +348,8 @@ def interp_decomp( :return: a tuple ``(k, idx, interp)`` containing the numerical rank *k*, the column indices *idx* and the resulting interpolation matrix *interp*. """ - if rank is not None and eps is not None: - raise ValueError("providing both 'rank' and 'eps' is not supported") + if (rank is not None and eps is not None) or (rank is None and eps is None): + raise ValueError("either 'rank' or 'eps' must be provided (not both)") import scipy.linalg.interpolative as sli if rank is None: @@ -486,3 +500,5 @@ def skeletonization_error( return result # }}} + +# vim: foldmethod=marker diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py index 78a7675c4..c24353dca 100644 --- a/pytential/symbolic/matrix.py +++ b/pytential/symbolic/matrix.py @@ -240,7 +240,7 @@ def map_call(self, expr): class ClusterMatrixBuilderBase(MatrixBuilderBase): """Evaluate individual clusters of a matrix operator, as defined by a - :class:`~pytential.lingla.TargetAndSourceClusterList`. + :class:`~pytential.linalg.utils.TargetAndSourceClusterList`. Unlike, e.g. :class:`MatrixBuilder`, matrix cluster builders are significantly reduced in scope. They are basically just meant @@ -252,7 +252,8 @@ class ClusterMatrixBuilderBase(MatrixBuilderBase): def __init__(self, actx, dep_expr, other_dep_exprs, dep_discr, places, tgt_src_index, context): """ - :arg tgt_src_index: a :class:`~pytential.linalg.TargetAndSourceClusterList` + :arg tgt_src_index: a + :class:`~pytential.linalg.utils.TargetAndSourceClusterList` class describing which clusters are going to be evaluated. """ diff --git a/test/extra_matrix_data.py b/test/extra_matrix_data.py index 022976073..398fce076 100644 --- a/test/extra_matrix_data.py +++ b/test/extra_matrix_data.py @@ -52,8 +52,8 @@ def get_cluster_index(self, actx, places, dofdesc=None): if max_particles_in_box is None: max_particles_in_box = discr.ndofs // self.approx_cluster_count - from pytential.linalg import partition_by_nodes - cindex = partition_by_nodes(actx, places, + from pytential.linalg.cluster import partition_by_nodes + cindex, ctree = partition_by_nodes(actx, places, dofdesc=dofdesc, tree_kind=self.tree_kind, max_particles_in_box=max_particles_in_box) @@ -74,12 +74,12 @@ def get_cluster_index(self, actx, places, dofdesc=None): from pytential.linalg import make_index_list cindex = make_index_list(subset) - return cindex + return cindex, ctree def get_tgt_src_cluster_index(self, actx, places, dofdesc=None): from pytential.linalg import TargetAndSourceClusterList - cindex = self.get_cluster_index(actx, places, dofdesc=dofdesc) - return TargetAndSourceClusterList(cindex, cindex) + cindex, ctree = self.get_cluster_index(actx, places, dofdesc=dofdesc) + return TargetAndSourceClusterList(cindex, cindex), ctree def get_operator(self, ambient_dim, qbx_forced_limit=_NoArgSentinel): knl = self.knl_class(ambient_dim) diff --git a/test/test_linalg_cluster.py b/test/test_linalg_cluster.py new file mode 100644 index 000000000..9cfb97f9b --- /dev/null +++ b/test/test_linalg_cluster.py @@ -0,0 +1,129 @@ +__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. +""" + +import pytest + +import numpy as np + +from pytential import GeometryCollection +from meshmode.mesh.generation import NArmedStarfish + +from meshmode import _acf # noqa: F401 +from arraycontext import pytest_generate_tests_for_array_contexts +from meshmode.array_context import PytestPyOpenCLArrayContextFactory + +import extra_matrix_data as extra +import logging +logger = logging.getLogger(__name__) + +pytest_generate_tests = pytest_generate_tests_for_array_contexts([ + PytestPyOpenCLArrayContextFactory, + ]) + +CLUSTER_TEST_CASES = [ + extra.CurveTestCase( + name="starfish", + target_order=4, + curve_fn=NArmedStarfish(5, 0.25), + resolutions=[64]), + extra.TorusTestCase( + target_order=4, + resolutions=[1]) + ] + + +# {{{ test_cluster_tree + +@pytest.mark.parametrize(("case", "tree_kind"), [ + (CLUSTER_TEST_CASES[0], None), + (CLUSTER_TEST_CASES[0], "adaptive"), + (CLUSTER_TEST_CASES[0], "adaptive-level-restricted"), + (CLUSTER_TEST_CASES[1], "adaptive"), + ]) +def test_cluster_tree(actx_factory, case, tree_kind, visualize=False): + if visualize: + logging.basicConfig(level=logging.INFO) + + from dataclasses import replace + actx = actx_factory() + case = replace(case, tree_kind=tree_kind) + logger.info("\n%s", case) + + discr = case.get_discretization(actx, case.resolutions[-1], case.target_order) + places = GeometryCollection(discr, auto_where=case.name) + + srcindex, ctree = case.get_cluster_index(actx, places) + assert srcindex.nclusters == ctree.nclusters + + from pytential.linalg.cluster import split_array + rng = np.random.default_rng(42) + x = split_array(rng.random(srcindex.indices.shape), srcindex) + + logger.info("nclusters %4d nlevels %4d", srcindex.nclusters, ctree.nlevels) + + if visualize and ctree._tree is not None: + import matplotlib.pyplot as plt + fig = plt.figure(figsize=(10, 10), dpi=300) + + from boxtree.visualization import TreePlotter + plotter = TreePlotter(ctree._tree) + plotter.draw_tree(fill=False, edgecolor="black", zorder=10) + plotter.draw_box_numbers() + plotter.set_bounding_box() + + fig.savefig("test_cluster_tree") + + from pytential.linalg.cluster import cluster, uncluster + for clevel in ctree.levels: + logger.info("======== Level %d", clevel.level) + logger.info("box_ids %s", clevel.box_ids) + logger.info("sizes %s", np.diff(srcindex.starts)) + logger.info("parent_map %s", clevel.parent_map) + + assert srcindex.nclusters == clevel.nclusters + + next_srcindex = cluster(srcindex, clevel) + for i, ppm in enumerate(clevel.parent_map): + partition = np.concatenate([srcindex.cluster_indices(j) for j in ppm]) + + assert partition.size == next_srcindex.cluster_size(i) + assert np.allclose(partition, next_srcindex.cluster_indices(i)) + + y = cluster(x, clevel) + z = uncluster(y, srcindex, clevel) + assert all(np.allclose(xi, zi) for xi, zi in zip(x, z)) + + srcindex = next_srcindex + x = y + +# }}} + + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: fdm=marker diff --git a/test/test_linalg_proxy.py b/test/test_linalg_proxy.py index 99029b63d..262bc0297 100644 --- a/test/test_linalg_proxy.py +++ b/test/test_linalg_proxy.py @@ -29,7 +29,7 @@ from arraycontext import flatten, unflatten from pytential import bind, sym from pytential import GeometryCollection -from pytential.linalg import ProxyGenerator, QBXProxyGenerator +from pytential.linalg.proxy import ProxyGenerator, QBXProxyGenerator from meshmode.mesh.generation import ellipse, NArmedStarfish from meshmode import _acf # noqa: F401 @@ -214,7 +214,7 @@ def test_partition_points(actx_factory, tree_kind, case, visualize=False): places = GeometryCollection(qbx, auto_where=case.name) density_discr = places.get_discretization(case.name) - mindex = case.get_cluster_index(actx, places) + mindex, _ = case.get_cluster_index(actx, places) expected_indices = np.arange(0, density_discr.ndofs) assert mindex.starts[-1] == density_discr.ndofs @@ -261,7 +261,7 @@ def test_proxy_generator(actx_factory, case, places = GeometryCollection(qbx, auto_where=case.name) density_discr = places.get_discretization(case.name) - cindex = case.get_cluster_index(actx, places) + cindex, _ = case.get_cluster_index(actx, places) generator = proxy_generator_cls(places, approx_nproxy=case.proxy_approx_count, @@ -330,7 +330,7 @@ def test_neighbor_points(actx_factory, case, dofdesc = places.auto_source density_discr = places.get_discretization(dofdesc.geometry) - srcindex = case.get_cluster_index(actx, places) + srcindex, _ = case.get_cluster_index(actx, places) # generate proxy points generator = proxy_generator_cls(places, @@ -339,7 +339,7 @@ def test_neighbor_points(actx_factory, case, pxy = generator(actx, dofdesc, srcindex) # get neighboring points - from pytential.linalg import gather_cluster_neighbor_points + from pytential.linalg.proxy import gather_cluster_neighbor_points nbrindex = gather_cluster_neighbor_points(actx, pxy) pxy = pxy.to_numpy(actx) diff --git a/test/test_linalg_skeletonization.py b/test/test_linalg_skeletonization.py index 2dcacbf04..b05cbeea9 100644 --- a/test/test_linalg_skeletonization.py +++ b/test/test_linalg_skeletonization.py @@ -118,37 +118,20 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False): places = GeometryCollection(qbx, auto_where=dd) density_discr = places.get_discretization(dd.geometry, dd.discr_stage) - tgt_src_index = case.get_tgt_src_cluster_index(actx, places, dd) + tgt_src_index, ctree = case.get_tgt_src_cluster_index(actx, places, dd) logger.info("nclusters %3d ndofs %7d", tgt_src_index.nclusters, density_discr.ndofs) # }}} - # {{{ wranglers - - from pytential.linalg import QBXProxyGenerator - from pytential.linalg.skeletonization import make_skeletonization_wrangler - proxy_generator = QBXProxyGenerator(places, - radius_factor=case.proxy_radius_factor, - approx_nproxy=case.proxy_approx_count) - + from pytential.linalg.skeletonization import rec_skeletonize_by_proxy sym_u, sym_op = case.get_operator(places.ambient_dim) - wrangler = make_skeletonization_wrangler(places, sym_op, sym_u, - domains=None, - context=case.knl_concrete_kwargs, - _weighted_proxy=case.weighted_proxy, - _proxy_source_cluster_builder=case.proxy_source_cluster_builder, - _proxy_target_cluster_builder=case.proxy_target_cluster_builder, - _neighbor_cluster_builder=case.neighbor_cluster_builder) - - # }}} - - from pytential.linalg.skeletonization import ( - _skeletonize_block_by_proxy_with_mats) - _skeletonize_block_by_proxy_with_mats( - actx, 0, 0, places, proxy_generator, wrangler, tgt_src_index, - id_eps=1.0e-8 + rec_skeletonize_by_proxy( + actx, places, ctree, tgt_src_index, sym_op, sym_u, + context=case.knl_concrete_kwargs, + auto_where=dd, + id_eps=1.0e-8, ) # }}} @@ -160,6 +143,7 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False): def run_skeletonize_by_proxy(actx, case, resolution, places=None, mat=None, ctol=None, rtol=None, + tgt_src_index=None, suffix="", visualize=False): from pytools import ProcessTimer @@ -169,9 +153,12 @@ def run_skeletonize_by_proxy(actx, case, resolution, if places is None: qbx = case.get_layer_potential(actx, resolution, case.target_order) places = GeometryCollection(qbx, auto_where=dd) + else: + qbx = places.get_geometry(dd.geometry) density_discr = places.get_discretization(dd.geometry, dd.discr_stage) - tgt_src_index = case.get_tgt_src_cluster_index(actx, places, dd) + if tgt_src_index is None: + tgt_src_index, _ = case.get_tgt_src_cluster_index(actx, places, dd) logger.info("nclusters %3d ndofs %7d", tgt_src_index.nclusters, density_discr.ndofs) @@ -188,7 +175,7 @@ def run_skeletonize_by_proxy(actx, case, resolution, logger.info("proxy factor %.2f count %7d", case.proxy_radius_factor, proxy_approx_count) - from pytential.linalg import QBXProxyGenerator + from pytential.linalg.proxy import QBXProxyGenerator from pytential.linalg.skeletonization import make_skeletonization_wrangler proxy_generator = QBXProxyGenerator(places, radius_factor=case.proxy_radius_factor, @@ -196,8 +183,8 @@ def run_skeletonize_by_proxy(actx, case, resolution, sym_u, sym_op = case.get_operator(places.ambient_dim) wrangler = make_skeletonization_wrangler(places, sym_op, sym_u, - domains=None, context=case.knl_concrete_kwargs, + auto_where=dd, _weighted_proxy=case.weighted_proxy, _proxy_source_cluster_builder=case.proxy_source_cluster_builder, _proxy_target_cluster_builder=case.proxy_target_cluster_builder, @@ -241,10 +228,13 @@ def run_skeletonize_by_proxy(actx, case, resolution, logger.info("[time] skeletonization by proxy: %s", p) + def intersect1d(x, y): + return np.where((x.reshape(1, -1) - y.reshape(-1, 1)) == 0)[1] + L, R = skeleton.L, skeleton.R for i in range(tgt_src_index.nclusters): # targets (rows) - bi = np.searchsorted( + bi = intersect1d( tgt_src_index.targets.cluster_indices(i), skeleton.skel_tgt_src_index.targets.cluster_indices(i), ) @@ -254,7 +244,7 @@ def run_skeletonize_by_proxy(actx, case, resolution, tgt_error = la.norm(A - L[i] @ S, ord=ord) / la.norm(A, ord=ord) # sources (columns) - bj = np.searchsorted( + bj = intersect1d( tgt_src_index.sources.cluster_indices(i), skeleton.skel_tgt_src_index.sources.cluster_indices(i), ) @@ -268,8 +258,8 @@ def run_skeletonize_by_proxy(actx, case, resolution, src_error, tgt_error, R[i].shape[0], R[i].shape[1]) if ctol is not None: - assert src_error < ctol * case.id_eps - assert tgt_error < ctol * case.id_eps + assert src_error < ctol + assert tgt_error < ctol # }}} @@ -300,9 +290,9 @@ def run_skeletonize_by_proxy(actx, case, resolution, rtol if rtol is not None else 0.0) if rtol: - assert err_l < rtol * case.id_eps - assert err_r < rtol * case.id_eps - assert err_f < rtol * case.id_eps + assert err_l < rtol + assert err_r < rtol + assert err_f < rtol # }}} @@ -342,19 +332,18 @@ def run_skeletonize_by_proxy(actx, case, resolution, # }}} - return err_f, (places, mat) + return err_f, (places, mat, skeleton) @pytest.mark.parametrize("case", [ - # NOTE: skip 2d tests, since they're better checked for convergence in - # `test_skeletonize_by_proxy_convergence` - # SKELETONIZE_TEST_CASES[0], SKELETONIZE_TEST_CASES[1], + SKELETONIZE_TEST_CASES[0], + SKELETONIZE_TEST_CASES[1], SKELETONIZE_TEST_CASES[2], ]) def test_skeletonize_by_proxy(actx_factory, case, visualize=False): - r"""Test single-level skeletonization accuracy. Checks that the error - satisfies :math:`e < c \epsilon_{id}` for a fixed ID tolerance and an - empirically determined (not too huge) :math:`c`. + r"""Test multilevel skeletonization accuracy. Checks that the error for + every level satisfies :math:`e < c \epsilon_{id}` for a fixed ID tolerance + and an empirically determined (not too huge) :math:`c`. """ import scipy.linalg.interpolative as sli @@ -368,12 +357,27 @@ def test_skeletonize_by_proxy(actx_factory, case, visualize=False): case = replace(case, approx_cluster_count=6, id_eps=1.0e-8) logger.info("\n%s", case) - run_skeletonize_by_proxy( - actx, case, case.resolutions[0], - ctol=6, - # FIXME: why is the 3D error so large? - rtol=10**case.ambient_dim, - visualize=visualize) + dd = sym.DOFDescriptor(case.name, discr_stage=case.skel_discr_stage) + qbx = case.get_layer_potential(actx, case.resolutions[0], case.target_order) + places = GeometryCollection(qbx, auto_where=dd) + + tgt_src_index, ctree = case.get_tgt_src_cluster_index(actx, places, dd) + mat = None + + from pytential.linalg.cluster import cluster + for clevel in ctree.levels[:-1]: + logger.info("[%2d/%2d] nclusters %3d", + clevel.level, ctree.nlevels, clevel.nclusters) + + _, (_, mat, skeleton) = run_skeletonize_by_proxy( + actx, case, case.resolutions[0], + ctol=6 * case.id_eps, + # FIXME: why is the 3D error so large? + rtol=10**case.ambient_dim * case.id_eps, + places=places, mat=mat, tgt_src_index=tgt_src_index, + visualize=visualize) + + tgt_src_index = cluster(skeleton.skel_tgt_src_index, clevel) # }}} @@ -437,9 +441,11 @@ def test_skeletonize_by_proxy_convergence( places = mat = None for i in range(id_eps.size): case = replace(case, id_eps=id_eps[i], weighted_proxy=weighted) - rec_error[i], (places, mat) = run_skeletonize_by_proxy( - actx, case, r, places=places, mat=mat, - suffix=f"{suffix}_{i:04d}", visualize=False) + + if not was_zero: + rec_error[i], (places, mat, _) = run_skeletonize_by_proxy( + actx, case, r, places=places, mat=mat, + suffix=f"{suffix}_{i:04d}", visualize=False) was_zero = rec_error[i] == 0.0 eoc.add_data_point(id_eps[i], rec_error[i]) diff --git a/test/test_matrix.py b/test/test_matrix.py index 3ab15ad47..22f159d7e 100644 --- a/test/test_matrix.py +++ b/test/test_matrix.py @@ -61,7 +61,7 @@ def max_cluster_error(mat, clusters, mindex, p=None): error = max( error, - la.norm(mat_i - clusters[i, i], ord=p) / norm_mat_i + la.norm(mat_i - clusters[i], ord=p) / norm_mat_i ) return error @@ -346,7 +346,7 @@ def test_cluster_builder(actx_factory, ambient_dim, # {{{ matrix - mindex = case.get_tgt_src_cluster_index(actx, places) + mindex, _ = case.get_tgt_src_cluster_index(actx, places) kwargs = { "dep_expr": sym_u, "other_dep_exprs": [], @@ -468,8 +468,8 @@ def test_build_matrix_fixed_stage(actx_factory, logger.info("ndofs: %d", target_discr.ndofs) from pytential.linalg import TargetAndSourceClusterList - itargets = case.get_cluster_index(actx, places, target_dd) - jsources = case.get_cluster_index(actx, places, source_dd) + itargets, _ = case.get_cluster_index(actx, places, target_dd) + jsources, _ = case.get_cluster_index(actx, places, source_dd) mindex = TargetAndSourceClusterList(itargets, jsources) kwargs = {