From 316479916fbdcf9741668246e9b08be709b5608a Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Thu, 4 Aug 2022 16:21:49 +0300 Subject: [PATCH] implement hierachical matrix solver --- pytential/linalg/cluster.py | 4 +- pytential/linalg/hmatrix.py | 211 ++++++++++++++++++++++++++++ pytential/linalg/skeletonization.py | 19 ++- pytential/linalg/utils.py | 9 ++ test/test_linalg_hmatrix.py | 116 +++++++++++++++ 5 files changed, 352 insertions(+), 7 deletions(-) create mode 100644 pytential/linalg/hmatrix.py create mode 100644 test/test_linalg_hmatrix.py diff --git a/pytential/linalg/cluster.py b/pytential/linalg/cluster.py index 1b085c569..e8b1b7d7a 100644 --- a/pytential/linalg/cluster.py +++ b/pytential/linalg/cluster.py @@ -129,7 +129,7 @@ class ClusterTree: def nclusters(self): return self.leaf_cluster_box_ids.size - def levels(self) -> Iterator[ClusterLevel]: + def levels(self, root: bool = False) -> Iterator[ClusterLevel]: """ :returns: an iterator over all the :class:`ClusterLevel` levels. """ @@ -142,7 +142,7 @@ def levels(self) -> Iterator[ClusterLevel]: parent_map=make_cluster_parent_map(parent_ids), ) - for _ in range(self.nlevels - 1, 0, -1): + for _ in range(self.nlevels - 1, -int(root), -1): yield clevel box_ids = np.unique(self.tree_cluster_parent_ids[clevel.box_ids]) diff --git a/pytential/linalg/hmatrix.py b/pytential/linalg/hmatrix.py new file mode 100644 index 000000000..2d2adc51d --- /dev/null +++ b/pytential/linalg/hmatrix.py @@ -0,0 +1,211 @@ +__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 +from typing import Any, Dict, Iterable, Optional, Union + +import numpy as np + +from arraycontext import PyOpenCLArrayContext, ArrayOrContainerT +from meshmode.dof_array import DOFArray + +from pytential import GeometryCollection, sym +from pytential.linalg.cluster import ClusterTree, cluster + +__doc__ = """ +Hierarical Matrix Construction +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +""" + + +# {{{ ProxyHierarchicalMatrix + +@dataclass(frozen=True) +class ProxyHierarchicalMatrix: + """ + .. attribute:: ctree + + A :class:`~pytential.linalg.cluster.ClusterTree`. + + .. attribute:: skeletons + + An :class:`~numpy.ndarray` containing skeletonization information + for each level of the hierarchy. For additional details, see + :class:`~pytential.linalg.skeletonization.SkeletonizationResult`. + + This class implements the :class:`scipy.sparse.linalg.LinearOperator` + interface. In particular, the following attributes and methods: + + .. attribute:: shape + + A :class:`tuple` that gives the matrix size ``(m, n)``. + + .. attribute:: dtype + + The data type of the matrix entries. + + .. automethod:: matvec + .. automethod:: __matmul__ + """ + + ctree: ClusterTree + skeletons: np.ndarray + + @property + def shape(self): + return self.skeletons[0].tgt_src_index.shape + + @property + def dtype(self): + # FIXME: assert that everyone has this dtype? + return self.skeletons[0].R[0].dtype + + @property + def nlevels(self): + return self.skeletons.size + + def matvec(self, x: ArrayOrContainerT) -> ArrayOrContainerT: + """Implements a matrix-vector multiplication :math:`H x`.""" + from arraycontext import get_container_context_recursively_opt + actx = get_container_context_recursively_opt(x) + if actx is None: + raise ValueError("input array is frozen") + + return apply_skeleton_matvec(actx, self, x) + + def __matmul__(self, x: ArrayOrContainerT) -> ArrayOrContainerT: + """Same as :meth:`matvec`.""" + return self.matvec(x) + + def rmatvec(self, x): + raise NotImplementedError + + def matmat(self, mat): + raise NotImplementedError + + def rmatmat(self, mat): + raise NotImplementedError + + +def apply_skeleton_matvec( + actx: PyOpenCLArrayContext, + hmat: ProxyHierarchicalMatrix, + x: ArrayOrContainerT, + ) -> ArrayOrContainerT: + from arraycontext import flatten + x = actx.to_numpy(flatten(x, actx, leaf_class=DOFArray)) + + from pytential.linalg.utils import split_array + y = split_array(x, hmat.skeletons[0].tgt_src_index.sources) + + assert x.dtype == hmat.dtype + assert x.shape == (hmat.shape[1],) + + d_dot_y = np.empty(hmat.nlevels, dtype=object) + r_dot_y = np.empty(hmat.nlevels, dtype=object) + + # recurse down + for k, clevel in enumerate(hmat.ctree.levels(root=True)): + skeleton = hmat.skeletons[k] + assert skeleton.tgt_src_index.shape[1] == sum(xi.size for xi in y) + + d_dot_y_k = np.empty(skeleton.nclusters, dtype=object) + r_dot_y_k = np.empty(skeleton.nclusters, dtype=object) + + for i in range(skeleton.nclusters): + r_dot_y_k[i] = skeleton.R[i] @ y[i] + d_dot_y_k[i] = skeleton.D[i] @ y[i] + + r_dot_y[k] = r_dot_y_k + d_dot_y[k] = d_dot_y_k + y = cluster(r_dot_y_k, clevel) + + # recurse up + for k, skeleton in reversed(list(enumerate(hmat.skeletons))): + r_dot_y_k = r_dot_y[k] + d_dot_y_k = d_dot_y[k] + + result = np.empty(skeleton.nclusters, dtype=object) + for i in range(skeleton.nclusters): + result[i] = skeleton.L[i] @ r_dot_y_k[i] + d_dot_y_k[i] + + from arraycontext import unflatten + return unflatten( + x, + actx.from_numpy(np.concatenate(result)), + actx) + +# }}} + + +# {{{ build_hmatrix_matvec_by_proxy + +def build_hmatrix_matvec_by_proxy( + actx: PyOpenCLArrayContext, + places: GeometryCollection, + exprs: Union[sym.Expression, Iterable[sym.Expression]], + input_exprs: Union[sym.Expression, Iterable[sym.Expression]], *, + domains: Optional[Iterable[sym.DOFDescriptorLike]] = None, + context: Optional[Dict[str, Any]] = None, + id_eps: float = 1.0e-8, + + # NOTE: these are dev variables and can disappear at any time! + # TODO: plugin in error model to get an estimate for: + # * how many points we want per cluster? + # * how many proxy points we want? + # * how far away should the proxy points be? + # based on id_eps. How many of these should be user tunable? + _tree_kind: Optional[str] = "adaptive-level-restricted", + _max_particles_in_box: Optional[int] = None, + + _id_rank: Optional[int] = None, + + _approx_nproxy: Optional[int] = None, + _proxy_radius_factor: Optional[float] = None, + _proxy_cls: Optional[type] = None, + ): + from pytential.linalg.cluster import partition_by_nodes + cluster_index, ctree = partition_by_nodes( + actx, places, + tree_kind=_tree_kind, + max_particles_in_box=_max_particles_in_box) + + from pytential.linalg.utils import TargetAndSourceClusterList + tgt_src_index = TargetAndSourceClusterList( + targets=cluster_index, sources=cluster_index) + + from pytential.linalg.skeletonization import rec_skeletonize_by_proxy + skeletons = rec_skeletonize_by_proxy( + actx, places, ctree, tgt_src_index, exprs, input_exprs, + domains=domains, + context=context, + id_eps=id_eps, + id_rank=_id_rank, + approx_nproxy=_approx_nproxy, + proxy_radius_factor=_proxy_radius_factor, + max_particles_in_box=_max_particles_in_box, + _proxy_cls=_proxy_cls, + ) + + return ProxyHierarchicalMatrix(ctree=ctree, skeletons=skeletons) + +# }}} diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py index 2ecafcdbf..ebc5e62dc 100644 --- a/pytential/linalg/skeletonization.py +++ b/pytential/linalg/skeletonization.py @@ -305,7 +305,8 @@ def evaluate_self( cls = self.neighbor_cluster_builder return self._evaluate_expr( actx, places, cls, tgt_src_index, self.exprs[ibrow], - idomain=ibcol, _weighted=self.weighted_sources) + idomain=ibcol, _weighted=self.weighted_sources, + ) # {{{ nearfield @@ -604,7 +605,8 @@ def _skeletonize_block_by_proxy_with_mats( wrangler: SkeletonizationWrangler, tgt_src_index: "TargetAndSourceClusterList", *, id_eps: Optional[float] = None, id_rank: Optional[int] = None, - max_particles_in_box: Optional[int] = None + max_particles_in_box: Optional[int] = None, + evaluate_diagonal: bool = False, ) -> "SkeletonizationResult": nclusters = tgt_src_index.nclusters if nclusters == 1: @@ -772,6 +774,9 @@ def __post_init__(self): if self.R.shape != shape: raise ValueError(f"'R' has shape {self.R.shape}, expected {shape}") + if self.D.shape != shape: + raise ValueError(f"'D' has shape {self.L.shape}, expected {shape}") + @property def nclusters(self): return self.tgt_src_index.nclusters @@ -792,7 +797,9 @@ def skeletonize_by_proxy( id_eps: Optional[float] = None, id_rank: Optional[int] = None, - max_particles_in_box: Optional[int] = None) -> np.ndarray: + max_particles_in_box: Optional[int] = None, + + _evaluate_diagonal: bool = False) -> np.ndarray: r"""Evaluate and skeletonize a symbolic expression using proxy-based methods. :arg tgt_src_index: a :class:`~pytential.linalg.utils.TargetAndSourceClusterList` @@ -832,7 +839,8 @@ def skeletonize_by_proxy( 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) + max_particles_in_box=max_particles_in_box, + evaluate_diagonal=_evaluate_diagonal) return skels @@ -954,7 +962,8 @@ def rec_skeletonize_by_proxy( skeleton = _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) + max_particles_in_box=max_particles_in_box, + ) skeleton = _fix_skeleton_diagonal( skeleton, skel_per_level[i - 1], parent_level) diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py index d831b04b9..173af8097 100644 --- a/pytential/linalg/utils.py +++ b/pytential/linalg/utils.py @@ -334,6 +334,15 @@ def make_flat_cluster_diag( return cluster_mat + +def split_array(x: np.ndarray, index: IndexList) -> np.ndarray: + 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) + ]) + # }}} diff --git a/test/test_linalg_hmatrix.py b/test/test_linalg_hmatrix.py new file mode 100644 index 000000000..8fcfc0071 --- /dev/null +++ b/test/test_linalg_hmatrix.py @@ -0,0 +1,116 @@ +__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 functools import partial +import pytest + +from pytential import sym +from pytential import GeometryCollection + +from meshmode.mesh.generation import ellipse, 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, + ]) + + +HMATRIX_TEST_CASES = [ + extra.CurveTestCase( + name="ellipse", + op_type="scalar", + target_order=7, + curve_fn=partial(ellipse, 3.0)), + extra.CurveTestCase( + name="starfish", + op_type="scalar", + target_order=4, + curve_fn=NArmedStarfish(5, 0.25), + resolutions=[32]), + extra.TorusTestCase( + target_order=4, + op_type="scalar", + resolutions=[0]) + ] + + +@pytest.mark.parametrize("case", HMATRIX_TEST_CASES) +def test_hmatrix_forward_matvec(actx_factory, case, visualize=False): + actx = actx_factory() + + if visualize: + logging.basicConfig(level=logging.INFO) + logger.info("\n%s", case) + + # {{{ geometry + + 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) + + density_discr = places.get_discretization(dd.geometry, dd.discr_stage) + 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) + + # }}} + + # {{{ construct hmatrix + + from pytential.linalg.hmatrix import build_hmatrix_matvec_by_proxy + sym_u, sym_op = case.get_operator(places.ambient_dim) + + hmat = build_hmatrix_matvec_by_proxy(actx, places, sym_op, sym_u, + domains=[dd], + context=case.knl_concrete_kwargs, + id_eps=case.id_eps, + _tree_kind=case.tree_kind, + _max_particles_in_box=case.max_particles_in_box, + _approx_nproxy=case.proxy_approx_count, + _proxy_radius_factor=case.proxy_radius_factor, + ) + + x = actx.thaw(density_discr.nodes()[0]) + r = hmat @ x + + assert r is not None + + # }}} + + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: fdm=marker