diff --git a/doc/conf.py b/doc/conf.py index 8656b31ac..60311d1cc 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -70,6 +70,8 @@ "cl_array.Array": "obj:pyopencl.array.Array", # pymbolic "ArithmeticExpression": "obj:pymbolic.ArithmeticExpression", + "ArithmeticExpressionContainerTc": + "obj:pymbolic.typing.ArithmeticExpressionContainerTc", "Expression": "obj:pymbolic.typing.Expression", "MultiVector": "obj:pymbolic.geometric_algebra.MultiVector", "Variable": "class:pymbolic.primitives.Variable", @@ -90,6 +92,7 @@ # boxtree "FromSepSmallerCrit": "obj:boxtree.traversal.FromSepSmallerCrit", "TimingResult": "class:boxtree.timing.TimingResult", + "Tree": "obj:boxtree.tree.Tree", "TreeKind": "obj:boxtree.tree_build.TreeKind", # sumpy "ExpansionBase": "class:sumpy.expansion.ExpansionBase", @@ -117,6 +120,7 @@ "pytential.symbolic.primitives.ExpressionNode": "class:pytential.symbolic.primitives.ExpressionNode", "sym.DOFDescriptor": "class:pytential.symbolic.dof_desc.DOFDescriptor", + "sym.DOFDescriptorLike": "obj:pytential.symbolic.dof_desc.DOFDescriptorLike", "sym.IntG": "class:pytential.symbolic.primitives.IntG", "sym.var": "obj:pytential.symbolic.primitives.var", } diff --git a/doc/linalg.rst b/doc/linalg.rst index de26425a0..97f12acbf 100644 --- a/doc/linalg.rst +++ b/doc/linalg.rst @@ -32,6 +32,7 @@ Low-level Functionality All the classes and routines in this module are experimental and the API can change at any point. +.. automodule:: pytential.linalg.cluster .. automodule:: pytential.linalg.proxy .. automodule:: pytential.linalg.skeletonization diff --git a/pytential/linalg/cluster.py b/pytential/linalg/cluster.py new file mode 100644 index 000000000..7ae2fb7bc --- /dev/null +++ b/pytential/linalg/cluster.py @@ -0,0 +1,595 @@ +from __future__ import annotations + + +__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 logging +import pathlib +from dataclasses import dataclass, replace +from functools import singledispatch +from typing import TYPE_CHECKING, Any + +import numpy as np + +from arraycontext import PyOpenCLArrayContext +from meshmode.discretization import Discretization +from pytools import log_process, memoize_method, obj_array + +from pytential import GeometryCollection, sym +from pytential.linalg.utils import IndexList, TargetAndSourceClusterList +from pytential.qbx import QBXLayerPotentialSource + + +if TYPE_CHECKING: + from collections.abc import Iterator + + import optype.numpy as onp + + from boxtree.tree import Tree + from boxtree.tree_build import TreeKind + + from pytential.linalg.proxy import ProxyGenerator + + +logger = logging.getLogger(__name__) + +__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: onp.Array1D[np.integer], + ) -> obj_array.ObjectArray1D[onp.Array1D[np.integer]]: + """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 obj_array.new_1d([ + 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`. + + .. autoattribute:: level + .. autoattribute:: box_ids + .. autoattribute:: parent_map + .. autoproperty:: nclusters + """ + + level: int + """Current level that is represented.""" + box_ids: onp.Array1D[np.integer] + """Box IDs on the current level.""" + parent_map: obj_array.ObjectArray1D[onp.Array1D[np.integer]] + """An object :class:`~numpy.ndarray` containing buckets of child indices, + i.e. ``parent_map[i]`` contains all the child indices that will cluster + into the same parent. Note that this indexing is local to this level + and is not related to the tree indexing stored by the :class:`ClusterTree`. + """ + + @property + def nclusters(self) -> int: + """Number of clusters on the current level (same as number of boxes + in :attr:`box_ids`). + """ + return self.box_ids.size + + +@dataclass(frozen=True) +class ClusterTree: + r"""Hierarchical cluster representation. + + .. autoattribute:: nlevels + .. autoattribute:: leaf_cluster_box_ids + .. autoattribute:: tree_cluster_parent_ids + + .. autoproperty:: nclusters + .. autoproperty:: levels + .. automethod:: iter_levels + """ + + nlevels: int + """Total number of levels in the tree.""" + leaf_cluster_box_ids: onp.Array1D[np.integer] + """Box IDs for each cluster on the leaf level of the tree.""" + tree_cluster_parent_ids: onp.Array1D[np.integer] + """Parent box IDs for :attr:`leaf_cluster_box_ids`.""" + + # NOTE: only here to allow easier debugging + testing + _tree: Tree | None + + @property + def nclusters(self) -> int: + """Number of clusters in the leaf level of the tree.""" + return self.leaf_cluster_box_ids.size + + @property + @memoize_method + def levels(self) -> obj_array.ObjectArray1D[ClusterLevel]: + r"""An :class:`~numpy.ndarray` of :class:`ClusterLevel`\ s.""" + return obj_array.new_1d(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: onp.Array1D[Any], + index: IndexList) -> obj_array.ObjectArray1D[onp.Array1D[Any]]: + """ + :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 + + return obj_array.new_1d([ + index.cluster_take(x, i) for i in range(index.nclusters) + ]) + + +@singledispatch +def cluster(obj: object, 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 = obj_array.new_1d([ + 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: obj_array.ObjectArray1D[onp.ArrayND[Any]], + clevel: ClusterLevel) -> obj_array.ObjectArray1D[onp.ArrayND[Any]]: + 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 obj_array.new_1d([ + np.concatenate([obj[i] for i in ppm]) for ppm in clevel.parent_map + ]) + elif ndim == 2: + return obj_array.new_1d([ + 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: obj_array.ObjectArray1D[onp.Array1D[Any]], + index: IndexList, + clevel: ClusterLevel) -> obj_array.ObjectArray1D[onp.Array1D[Any]]: + """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: onp.Array1D[np.integer]) -> ClusterTree: + partition_box_ids = np.arange(starts.size - 1) + box_ids = partition_box_ids + + box_parent_ids: list[onp.Array1D[np.integer]] = [] + 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) + + +@log_process(logger) +def partition_by_nodes( + actx: PyOpenCLArrayContext, places: GeometryCollection, *, + dofdesc: sym.DOFDescriptorLike | None = None, + tree_kind: TreeKind | None = "adaptive-level-restricted", + max_particles_in_box: int | None = 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, Discretization | QBXLayerPotentialSource) + + discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage) + assert isinstance(discr, Discretization) + + if tree_kind is not None: + setup_actx = lpot_source._setup_actx + assert isinstance(setup_actx, PyOpenCLArrayContext) + + from pytential.qbx.utils import tree_code_container + tcc = tree_code_container(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) + tree = tree.get(actx.queue) + + # FIXME maybe this should use IS_LEAF once available? + from boxtree import box_flags_enum + assert tree.box_flags is not None + leaf_boxes, = ( + tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0 + ).nonzero() + + # FIXME: this annotation is not needed with numpy 2.0 + 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) + + # FIXME: mypy seems to be able to figure this out with numpy 2.0 + assert starts is not None + 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 + +# }}} + + +# {{{ visualize clusters + +def visualize_clusters(actx: PyOpenCLArrayContext, + generator: ProxyGenerator, + srcindex: IndexList, + tree: ClusterTree, + filename: str | pathlib.Path, *, + dofdesc: sym.DOFDescriptorLike = None, + overwrite: bool = False) -> None: + filename = pathlib.Path(filename) + + places = generator.places + if dofdesc is None: + dofdesc = places.auto_source + dofdesc = sym.as_dofdesc(dofdesc) + + discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage) + assert isinstance(discr, Discretization) + + if discr.ambient_dim == 2: + _visualize_clusters_2d(actx, generator, discr, srcindex, tree, filename, + dofdesc=dofdesc, overwrite=overwrite) + elif discr.ambient_dim == 2: + _visualize_clusters_3d(actx, generator, discr, srcindex, tree, filename, + dofdesc=dofdesc, overwrite=overwrite) + else: + raise NotImplementedError(f"Unsupported dimension: {discr.ambient_dim}") + + +def _visualize_clusters_2d(actx: PyOpenCLArrayContext, + generator: ProxyGenerator, + discr: Discretization, + srcindex: IndexList, + tree: ClusterTree, + filename: pathlib.Path, *, + dofdesc: sym.DOFDescriptor, + overwrite: bool = False) -> None: + import matplotlib.pyplot as pt + + from arraycontext import flatten + from boxtree.visualization import TreePlotter + from meshmode.dof_array import DOFArray + + assert discr.ambient_dim == 2 + x, y = actx.to_numpy(flatten(discr.nodes(), actx, leaf_class=DOFArray)) + for clevel in tree.levels: + outfile = filename.with_stem(f"{filename.stem}-{clevel.level:03d}") + if not overwrite and outfile.exists(): + raise FileExistsError(f"Output file '{outfile}' already exists") + + pxy = generator(actx, dofdesc, srcindex).to_numpy(actx) + pxycenters = pxy.centers + pxyradii = pxy.radii + clsradii = pxy.cluster_radii + + fig = pt.figure() + ax = fig.gca() + + plotter = TreePlotter(tree._tree) + plotter.set_bounding_box() + plotter.draw_tree(fill=False, edgecolor="black", zorder=10) + + ax.plot(x, y, "ko", ms=2.0) + for i in range(srcindex.nclusters): + isrc = srcindex.cluster_indices(i) + ax.plot(x[isrc], y[isrc], "o", ms=2.0) + + from itertools import cycle + colors = cycle(pt.rcParams["axes.prop_cycle"].by_key()["color"]) + + for ppm in clevel.parent_map: + color = next(colors) + for j in ppm: + center = (pxycenters[0, j], pxycenters[1, j]) + c = pt.Circle(center, pxyradii[j], color=color, alpha=0.1) + ax.add_artist(c) + c = pt.Circle(center, clsradii[j], color=color, alpha=0.1) + ax.add_artist(c) + ax.text(*center, f"{j}", fontsize=18) + + ax.set_xlabel("$x$") + ax.set_ylabel("$y$") + ax.relim() + ax.autoscale() + ax.set_aspect("equal") + + fig.savefig(outfile) + pt.close(fig) + + srcindex = cluster(srcindex, clevel) + + +def _visualize_clusters_3d(actx: PyOpenCLArrayContext, + generator: ProxyGenerator, + discr: Discretization, + srcindex: IndexList, + tree: ClusterTree, + filename: pathlib.Path, *, + dofdesc: sym.DOFDescriptor, + overwrite: bool = False) -> None: + from arraycontext import unflatten + from meshmode.discretization.visualization import make_visualizer + + # NOTE: This writes out one vtu file for each level that contains + # * a mesh that's the union of `discr` and a sphere for each proxy ball + # * marker: a marker on `discr` (NaN on the proxy balls) for each of the + # clusters at the current level + # * proxies: a marker on the proxy balls (NaN on `discr`) + # + # Not quite sure how to best visualize the whole geometry here, so the + # proposed workflow is to load the vtu file twice, set opacity to 0 for + # NaNs and set opacity to something small for the proxy balls. + + # TODO: + # * color proxy balls based on their parent so we can easily see how they + # will cluster + + assert discr.ambient_dim == 3 + for clevel in tree.levels: + outfile = filename.with_stem(f"{filename.stem}-lvl{clevel.level:03d}") + outfile = outfile.with_suffix(".vtu") + if not overwrite and outfile.exists(): + raise FileExistsError(f"Output file '{outfile}' already exists") + + # construct proxy balls + pxy = generator(actx, dofdesc, srcindex).to_numpy(actx) + pxycenters = pxy.centers + pxyradii = pxy.radii + nclusters = srcindex.nclusters + + # construct meshes for each proxy ball + from meshmode.mesh.generation import generate_sphere + from meshmode.mesh.processing import affine_map, merge_disjoint_meshes + + ref_mesh = generate_sphere(1, 4, uniform_refinement_rounds=1) + pxymeshes = [ + affine_map(ref_mesh, A=pxyradii[i], b=pxycenters[:, i].squeeze()) + for i in range(nclusters) + ] + + # merge meshes into a single discretization + from meshmode.discretization.poly_element import ( + InterpolatoryEdgeClusteredGroupFactory, + ) + pxymesh = merge_disjoint_meshes([discr.mesh, *pxymeshes]) + pxydiscr = Discretization(actx, pxymesh, + InterpolatoryEdgeClusteredGroupFactory(4)) + + # add a marker field for all clusters + marker = np.full((pxydiscr.ndofs,), np.nan, dtype=np.float64) + template_ary = actx.thaw(pxydiscr.nodes()[0]) + + for i in range(srcindex.nclusters): + isrc = srcindex.cluster_indices(i) + marker[isrc] = 10.0 * (i + 1.0) + marker_dev = unflatten(template_ary, actx.from_numpy(marker), actx) + + # add a marker field for all proxies + pxymarker = np.full((pxydiscr.ndofs,), np.nan, dtype=np.float64) + pxymarker[discr.ndofs:] = 1.0 + pxymarker_dev = unflatten(template_ary, actx.from_numpy(pxymarker), actx) + + # write it all out + vis = make_visualizer(actx, pxydiscr) + vis.write_vtk_file(str(outfile), [ + ("marker", marker_dev), + ("proxies", pxymarker_dev), + ], overwrite=overwrite) + + srcindex = cluster(srcindex, clevel) + + +# }}} + + +# vim: foldmethod=marker diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py index 964698a8e..cebaba70e 100644 --- a/pytential/linalg/proxy.py +++ b/pytential/linalg/proxy.py @@ -36,7 +36,7 @@ from arraycontext import Array, ArrayContainer, PyOpenCLArrayContext, flatten from meshmode.discretization import Discretization from meshmode.dof_array import DOFArray -from pytools import memoize_in +from pytools import log_process, memoize_in from pytential import GeometryCollection, bind, sym from pytential.qbx import QBXLayerPotentialSource @@ -49,7 +49,6 @@ import optype.numpy as onp - from boxtree.tree_build import TreeKind from sumpy.expansion import ExpansionBase from sumpy.kernel import Kernel @@ -59,7 +58,6 @@ logger = logging.getLogger(__name__) - __doc__ = """ Proxy Point Generation ~~~~~~~~~~~~~~~~~~~~~~ @@ -74,7 +72,6 @@ .. autoclass:: QBXProxyGenerator :show-inheritance: -.. autofunction:: partition_by_nodes .. autofunction:: gather_cluster_neighbor_points """ @@ -82,82 +79,6 @@ _DEFAULT_MAX_PARTICLES_IN_BOX = 32 -# {{{ point index partitioning - -def partition_by_nodes( - actx: PyOpenCLArrayContext, - places: GeometryCollection, *, - dofdesc: DOFDescriptorLike | None = None, - tree_kind: TreeKind | None = "adaptive-level-restricted", - max_particles_in_box: int | None = 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? - assert tree.box_flags is not None - leaf_boxes, = ( - tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0 - ).nonzero() - - indices = np.empty(len(leaf_boxes), dtype=object) - starts: onp.Array1D[np.integer] | None = 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.utils import make_index_list - return make_index_list(indices, starts=starts) - -# }}} - - # {{{ proxy points class ProxyPointSource(PointPotentialSource): @@ -469,6 +390,7 @@ def get_centers_kernel_ex(self, actx: PyOpenCLArrayContext) -> lp.ExecutorBase: def get_radii_kernel_ex(self, actx: PyOpenCLArrayContext) -> lp.ExecutorBase: pass + @log_process(logger) def __call__(self, actx: PyOpenCLArrayContext, source_dd: DOFDescriptorLike | None, @@ -655,6 +577,7 @@ def get_radii_kernel_ex(self, actx: PyOpenCLArrayContext) -> lp.ExecutorBase: return make_compute_cluster_qbx_radii_kernel_ex(actx, self.ambient_dim) @override + @log_process(logger) def __call__(self, actx: PyOpenCLArrayContext, source_dd: DOFDescriptorLike | None, @@ -688,6 +611,7 @@ def __call__(self, # {{{ gather_cluster_neighbor_points +@log_process(logger) def gather_cluster_neighbor_points( actx: PyOpenCLArrayContext, pxy: ProxyClusterGeometryData, diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py index 0f14fe5dc..5f7b05033 100644 --- a/pytential/linalg/skeletonization.py +++ b/pytential/linalg/skeletonization.py @@ -23,15 +23,17 @@ THE SOFTWARE. """ +import logging from dataclasses import dataclass from typing import TYPE_CHECKING, Any import numpy as np from meshmode.discretization import Discretization -from pytools import memoize_in, obj_array +from pytools import log_process, memoize_in, obj_array from pytential import GeometryCollection, bind, sym +from pytential.linalg.cluster import ClusterTree, cluster from pytential.linalg.direct_solver_symbolic import ( PROXY_SKELETONIZATION_SOURCE, PROXY_SKELETONIZATION_TARGET, @@ -53,6 +55,8 @@ from pytential.symbolic.matrix import ClusterMatrixBuilderBase +logger = logging.getLogger(__name__) + __doc__ = """ Skeletonization --------------- @@ -62,6 +66,7 @@ .. autoclass:: SkeletonizationResult .. autofunction:: skeletonize_by_proxy +.. autofunction:: rec_skeletonize_by_proxy """ @@ -143,7 +148,9 @@ def prg(): """ <> ioffset = starts[icluster] <> npoints = starts[icluster + 1] - ioffset - result[icluster] = reduce(sum, i, waa[indices[i + ioffset]]) / npoints + result[icluster] = ( + reduce(sum, i, abs(waa[indices[i + ioffset]])) / npoints + if npoints > 0 else 1.0) """, lang_version=lp.MOST_RECENT_LANGUAGE_VERSION, ) @@ -307,8 +314,20 @@ def _evaluate_expr( context=self.context, **kwargs)(expr) + def evaluate_self(self, + actx: PyOpenCLArrayContext, + places: GeometryCollection, + tgt_src_index: TargetAndSourceClusterList, + ibrow: int, ibcol: int, + ) -> onp.Array1D[Any]: + cls = self.neighbor_cluster_builder + return self._evaluate_expr( + actx, places, cls, tgt_src_index, self.exprs[ibrow], + idomain=ibcol, _weighted=True) + # {{{ nearfield + @log_process(logger) def evaluate_source_neighbor_interaction(self, actx: PyOpenCLArrayContext, places: GeometryCollection, @@ -322,10 +341,11 @@ def evaluate_source_neighbor_interaction(self, expr = self.exprs[ibrow] mat = self._evaluate_expr( actx, places, eval_mapper_cls, nbr_src_index, expr, - idomain=ibcol, _weighted=self.weighted_sources) + idomain=ibcol, _weighted=True) return mat, nbr_src_index + @log_process(logger) def evaluate_target_neighbor_interaction(self, actx: PyOpenCLArrayContext, places: GeometryCollection, @@ -339,7 +359,7 @@ def evaluate_target_neighbor_interaction(self, expr = self.exprs[ibrow] mat = self._evaluate_expr( actx, places, eval_mapper_cls, tgt_nbr_index, expr, - idomain=ibcol, _weighted=self.weighted_targets) + idomain=ibcol, _weighted=True) return mat, tgt_nbr_index @@ -347,6 +367,7 @@ def evaluate_target_neighbor_interaction(self, # {{{ proxy + @log_process(logger) def evaluate_source_proxy_interaction(self, actx: PyOpenCLArrayContext, places: GeometryCollection, @@ -360,6 +381,10 @@ def evaluate_source_proxy_interaction(self, places, {PROXY_SKELETONIZATION_TARGET: pxy.as_targets()} ) + if not self.weighted_sources: + logger.warning("Source-Proxy weighting is turned off. This will not give " + "good results for skeletonization.", stacklevel=3) + eval_mapper_cls = self.proxy_source_cluster_builder expr = self.source_proxy_exprs[ibrow] mat = self._evaluate_expr( @@ -370,6 +395,7 @@ def evaluate_source_proxy_interaction(self, return mat, pxy_src_index + @log_process(logger) def evaluate_target_proxy_interaction(self, actx: PyOpenCLArrayContext, places: GeometryCollection, @@ -394,6 +420,9 @@ def evaluate_target_proxy_interaction(self, mat = _apply_weights( actx, mat, places, tgt_pxy_index, nbrindex, self.domains[ibcol]) + else: + logger.warning("Target-Proxy weighting is turned off. This will not give " + "good results for skeletonization.", stacklevel=3) return mat, tgt_pxy_index @@ -562,7 +591,8 @@ def _evaluate_proxy_skeletonization_interaction( actx: PyOpenCLArrayContext, places: GeometryCollection, proxy_generator: ProxyGeneratorBase, - cluster_index: IndexList, *, + source_index: IndexList, + target_index: IndexList, *, evaluate_proxy: Callable[..., tuple[onp.Array1D[np.inexact], TargetAndSourceClusterList]], evaluate_neighbor: Callable[..., @@ -574,23 +604,24 @@ def _evaluate_proxy_skeletonization_interaction( each cluster in *cluster_index*. """ - if cluster_index.nclusters == 1: + if source_index.nclusters == 1: raise ValueError("cannot make a proxy skeleton for a single cluster") from pytential.linalg.proxy import gather_cluster_neighbor_points - pxy = proxy_generator(actx, dofdesc, cluster_index) + pxy = proxy_generator(actx, dofdesc, source_index) nbrindex = gather_cluster_neighbor_points( - actx, pxy, + actx, pxy, target_index, max_particles_in_box=max_particles_in_box) pxymat, pxy_cluster_index = evaluate_proxy(actx, places, pxy, nbrindex) nbrmat, nbr_cluster_index = evaluate_neighbor(actx, places, pxy, nbrindex) - - return _ProxyNeighborEvaluationResult( + result = _ProxyNeighborEvaluationResult( pxy=pxy, pxymat=pxymat, pxyindex=pxy_cluster_index, nbrmat=nbrmat, nbrindex=nbr_cluster_index) + return result + def _skeletonize_block_by_proxy_with_mats( actx: PyOpenCLArrayContext, ibrow: int, ibcol: int, @@ -615,80 +646,109 @@ def _skeletonize_block_by_proxy_with_mats( dofdesc=wrangler.domains[ibcol], max_particles_in_box=max_particles_in_box) - src_result = evaluate_skeletonization_interaction( - tgt_src_index.sources, - evaluate_proxy=partial( - wrangler.evaluate_source_proxy_interaction, - ibrow=ibrow, ibcol=ibcol), - evaluate_neighbor=partial( - wrangler.evaluate_source_neighbor_interaction, - ibrow=ibrow, ibcol=ibcol), - ) - tgt_result = evaluate_skeletonization_interaction( - tgt_src_index.targets, - evaluate_proxy=partial( - wrangler.evaluate_target_proxy_interaction, - ibrow=ibrow, ibcol=ibcol), - evaluate_neighbor=partial( - wrangler.evaluate_target_neighbor_interaction, - ibrow=ibrow, ibcol=ibcol) - ) - - src_skl_indices = np.empty(nclusters, dtype=object) - tgt_skl_indices = np.empty(nclusters, dtype=object) + from pytential.linalg.utils import interp_decomp + skel_src_indices = np.empty(nclusters, dtype=object) + skel_tgt_indices = np.empty(nclusters, dtype=object) skel_starts = np.zeros(nclusters + 1, dtype=np.int32) L = np.empty(nclusters, dtype=object) R = np.empty(nclusters, dtype=object) - from pytential.linalg.utils import interp_decomp + from pytools import ProcessLogger + + with ProcessLogger( + logger, + f"_skeletonize_block_by_proxy_with_mats_{ibrow}_{ibcol}"): + src_result = evaluate_skeletonization_interaction( + tgt_src_index.sources, tgt_src_index.targets, + evaluate_proxy=partial( + wrangler.evaluate_source_proxy_interaction, + ibrow=ibrow, ibcol=ibcol), + evaluate_neighbor=partial( + wrangler.evaluate_source_neighbor_interaction, + ibrow=ibrow, ibcol=ibcol), + ) + tgt_result = evaluate_skeletonization_interaction( + tgt_src_index.targets, tgt_src_index.sources, + evaluate_proxy=partial( + wrangler.evaluate_target_proxy_interaction, + ibrow=ibrow, ibcol=ibcol), + evaluate_neighbor=partial( + wrangler.evaluate_target_neighbor_interaction, + ibrow=ibrow, ibcol=ibcol) + ) - for i in range(nclusters): - k = id_rank - src_mat = np.vstack(src_result[i]) - tgt_mat = np.hstack(tgt_result[i]) - max_allowable_rank = min(*src_mat.shape, *tgt_mat.shape) + for i in range(nclusters): + k = id_rank + src_mat = np.vstack(src_result[i]) + tgt_mat = np.hstack(tgt_result[i]) + max_allowable_rank = min(*src_mat.shape, *tgt_mat.shape) - if __debug__: - isfinite = np.isfinite(tgt_mat) - assert np.all(isfinite), np.where(isfinite) - isfinite = np.isfinite(src_mat) - assert np.all(isfinite), np.where(isfinite) + if __debug__: + isfinite = np.isfinite(tgt_mat) + assert np.all(isfinite), np.where(isfinite) + isfinite = np.isfinite(src_mat) + assert np.all(isfinite), np.where(isfinite) - # skeletonize target points - k, idx, interp = interp_decomp(tgt_mat.T, rank=k, eps=id_eps, rng=rng) - assert 0 < k <= len(idx) + # skeletonize target points + k, idx, interp = interp_decomp(tgt_mat.T, rank=k, eps=id_eps, rng=rng) + assert 0 < k <= len(idx) - if k > max_allowable_rank: - k = max_allowable_rank - interp = interp[:k, :] + if k > max_allowable_rank: + k = max_allowable_rank + interp = interp[:k, :] - L[i] = interp.T - tgt_skl_indices[i] = tgt_src_index.targets.cluster_indices(i)[idx[:k]] - assert interp.shape == (k, tgt_mat.shape[0]) + L[i] = interp.T + skel_tgt_indices[i] = tgt_src_index.targets.cluster_indices(i)[idx[:k]] + assert interp.shape == (k, tgt_mat.shape[0]) - # skeletonize source points - k, idx, interp = interp_decomp(src_mat, rank=k, eps=None, rng=rng) - assert 0 < k <= len(idx) + # skeletonize source points + k, idx, interp = interp_decomp(src_mat, rank=k, eps=None, rng=rng) + assert 0 < k <= len(idx) - R[i] = interp - src_skl_indices[i] = tgt_src_index.sources.cluster_indices(i)[idx[:k]] - assert interp.shape == (k, src_mat.shape[1]) + R[i] = interp + skel_src_indices[i] = tgt_src_index.sources.cluster_indices(i)[idx[:k]] + assert interp.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 + skel_starts[i + 1] = skel_starts[i] + k + assert skel_tgt_indices[i].shape == skel_src_indices[i].shape - from pytential.linalg.utils import make_index_list + # 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) - 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) + from pytential.linalg import 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) + + return SkeletonizationResult( + L=obj_array.new_1d([np.eye(*D[0].shape)]), + R=obj_array.new_1d([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, + ) + # }}} @@ -719,6 +779,7 @@ class SkeletonizationResult: .. autoattribute:: L .. autoattribute:: R + .. autoattribute:: D .. autoattribute:: tgt_src_index .. autoattribute:: skel_tgt_src_index """ @@ -729,6 +790,9 @@ class SkeletonizationResult: R: obj_array.ObjectArray1D[onp.Array2D[Any]] """An object :class:`~numpy.ndarray` of size ``(nclusters,)`` that contains the right block interpolation matrices.""" + D: obj_array.ObjectArray1D[onp.Array2D[Any]] + """An object :class:`~numpy.ndarray` of size ``(nclusters,)`` that contains + the dense diagonal blocks.""" tgt_src_index: TargetAndSourceClusterList """A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` representing @@ -767,15 +831,17 @@ def nclusters(self) -> int: return self.tgt_src_index.nclusters +@log_process(logger) def skeletonize_by_proxy( actx: PyOpenCLArrayContext, places: GeometryCollection, tgt_src_index: TargetAndSourceClusterList, - exprs: sym.var | Sequence[sym.var], + exprs: ArithmeticExpression | Sequence[ArithmeticExpression], input_exprs: sym.var | Sequence[sym.var], *, domains: Sequence[Hashable] | None = None, context: dict[str, Any] | None = None, + auto_where: Any = None, approx_nproxy: int | None = None, proxy_radius_factor: float | None = None, @@ -784,7 +850,7 @@ def skeletonize_by_proxy( id_rank: int | None = None, rng: np.random.Generator | None = None, max_particles_in_box: int | None = None, - ) -> obj_array.ObjectArray2D[onp.Array2D[np.inexact]]: + ) -> obj_array.ObjectArray2D[SkeletonizationResult]: r"""Evaluate and skeletonize a symbolic expression using proxy-based methods. :arg tgt_src_index: a :class:`~pytential.linalg.utils.TargetAndSourceClusterList` @@ -810,21 +876,105 @@ 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) + from itertools import product + skels = 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, - rng=rng) + for ibrow, ibcol in product(range(wrangler.nrows), 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, + rng=rng) return skels # }}} + + +# {{{ recursive skeletonization by proxy + +@log_process(logger) +def rec_skeletonize_by_proxy( + actx: PyOpenCLArrayContext, + places: GeometryCollection, + + ctree: ClusterTree, + tgt_src_index: TargetAndSourceClusterList, + exprs: ArithmeticExpression | Sequence[ArithmeticExpression], + input_exprs: sym.var | Sequence[sym.var], *, + domains: Sequence[Hashable] | None = None, + context: dict[str, Any] | None = None, + auto_where: Any = None, + + approx_nproxy: int | None = None, + proxy_radius_factor: float | None = None, + + id_eps: float | None = None, + rng: np.random.Generator | None = None, + max_particles_in_box: int | None = None, + + _wrangler: SkeletonizationWrangler | None = None, + _proxy: ProxyGeneratorBase | None = None, + ) -> obj_array.ObjectArray1D[SkeletonizationResult]: + 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.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, + rng=rng, + 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 ed86aff85..aac33789f 100644 --- a/pytential/linalg/utils.py +++ b/pytential/linalg/utils.py @@ -53,6 +53,9 @@ .. autofunction:: make_index_list .. autofunction:: make_index_cluster_cartesian_product +.. autofunction:: make_flat_cluster_diag + +.. autofunction:: interp_decomp """ InexactT = TypeVar("InexactT", bound=np.inexact) diff --git a/test/extra_matrix_data.py b/test/extra_matrix_data.py index 46ff9686f..544ad171b 100644 --- a/test/extra_matrix_data.py +++ b/test/extra_matrix_data.py @@ -58,8 +58,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.proxy 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) @@ -81,12 +81,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..25605993f --- /dev/null +++ b/test/test_linalg_cluster.py @@ -0,0 +1,133 @@ +from __future__ import annotations + + +__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 logging + +import extra_matrix_data as extra +import numpy as np +import pytest + +from arraycontext import pytest_generate_tests_for_array_contexts +from meshmode import _acf # noqa: F401 +from meshmode.array_context import PytestPyOpenCLArrayContextFactory +from meshmode.mesh.generation import NArmedStarfish + +from pytential import GeometryCollection + + +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, strict=True)) + + 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 9115cdfd4..a6606f526 100644 --- a/test/test_linalg_proxy.py +++ b/test/test_linalg_proxy.py @@ -212,7 +212,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 @@ -259,7 +259,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, @@ -305,7 +305,7 @@ def test_proxy_generator(actx_factory, case, ProxyGenerator, QBXProxyGenerator, ]) @pytest.mark.parametrize("index_sparsity_factor", [1.0, 0.6]) -@pytest.mark.parametrize("proxy_radius_factor", [1, 1.1]) +@pytest.mark.parametrize("proxy_radius_factor", [1.0, 1.1]) def test_neighbor_points(actx_factory, case, proxy_generator_cls, index_sparsity_factor, proxy_radius_factor, visualize=False): @@ -328,7 +328,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, @@ -338,7 +338,7 @@ def test_neighbor_points(actx_factory, case, # get neighboring points from pytential.linalg.proxy import gather_cluster_neighbor_points - nbrindex = gather_cluster_neighbor_points(actx, pxy) + nbrindex = gather_cluster_neighbor_points(actx, pxy, srcindex) pxy = pxy.to_numpy(actx) nodes = actx.to_numpy( diff --git a/test/test_linalg_skeletonization.py b/test/test_linalg_skeletonization.py index 675b69362..619fed7a1 100644 --- a/test/test_linalg_skeletonization.py +++ b/test/test_linalg_skeletonization.py @@ -127,36 +127,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.skeletonization import rec_skeletonize_by_proxy - from pytential.linalg.proxy import QBXProxyGenerator - proxy_generator = QBXProxyGenerator(places, - radius_factor=case.proxy_radius_factor, - approx_nproxy=case.proxy_approx_count) - - from pytential.linalg.skeletonization import make_skeletonization_wrangler 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, + 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, rng=rng, ) @@ -186,7 +170,7 @@ def run_skeletonize_by_proxy(actx, case, resolution, density_discr = places.get_discretization(dd.geometry, dd.discr_stage) if tgt_src_index is None: - tgt_src_index = case.get_tgt_src_cluster_index(actx, places, dd) + 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) @@ -362,19 +346,18 @@ def intersect1d(x, y): # }}} - 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 @@ -390,13 +373,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=10 * case.id_eps, - # FIXME: why is the 3D error so large? - rtol=10**case.ambient_dim * case.id_eps, - rng=rng, - 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=10 * case.id_eps, + # FIXME: why is the 3D error so large? + rtol=10**case.ambient_dim * case.id_eps, + places=places, mat=mat, rng=rng, tgt_src_index=tgt_src_index, + visualize=visualize) + + tgt_src_index = cluster(skeleton.skel_tgt_src_index, clevel) # }}} @@ -466,7 +463,7 @@ def test_skeletonize_by_proxy_convergence( # NOTE: don't skeletonize anymore if we reached zero error, but we still # want to loop to do `eoc.add_data_point()` if not was_zero: - rec_error[i], (places, mat) = run_skeletonize_by_proxy( + rec_error[i], (places, mat, _) = run_skeletonize_by_proxy( actx, case, r, places=places, mat=mat, suffix=f"{suffix}_{i:04d}", rng=rng, visualize=False) diff --git a/test/test_matrix.py b/test/test_matrix.py index 97f27d38c..8f849a882 100644 --- a/test/test_matrix.py +++ b/test/test_matrix.py @@ -348,7 +348,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": [], @@ -472,8 +472,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 = {