diff --git a/pytential/linalg/cluster.py b/pytential/linalg/cluster.py index bae4911f6..1e58d1cc7 100644 --- a/pytential/linalg/cluster.py +++ b/pytential/linalg/cluster.py @@ -220,6 +220,23 @@ def make_block(i: int, j: int): else: raise ValueError(f"unsupported ndarray dimension: '{ndim}'") + +def uncluster(ary: np.ndarray, index: IndexList, clevel: ClusterLevel) -> np.ndarray: + assert ary.shape == (clevel.parent_map.size,) + + if index.nclusters == 1: + return ary + + result = 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 + + return result + # }}} diff --git a/pytential/linalg/hmatrix.py b/pytential/linalg/hmatrix.py new file mode 100644 index 000000000..174f523e0 --- /dev/null +++ b/pytential/linalg/hmatrix.py @@ -0,0 +1,344 @@ +__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 +import numpy.linalg as la + +from arraycontext import PyOpenCLArrayContext, ArrayOrContainerT, flatten, unflatten +from meshmode.dof_array import DOFArray + +from pytential import GeometryCollection, sym +from pytential.linalg.cluster import ClusterTree +from pytential.linalg.skeletonization import SkeletonizationWrangler + +__doc__ = """ +Hierarical Matrix Construction +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +""" + + +# {{{ ProxyHierarchicalMatrix + +@dataclass(frozen=True) +class ProxyHierarchicalMatrix: + """ + .. 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__ + """ + + wrangler: SkeletonizationWrangler + ctree: ClusterTree + method: str + + 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 + + @property + def nclusters(self): + return self.skeletons[0].nclusters + + 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") + + if self.method == "forward": + return apply_skeleton_matvec(actx, self, x) + elif self.method == "backward": + return apply_skeleton_inverse_matvec(actx, self, x) + else: + raise ValueError(f"unknown matvec method: '{self.method}'") + + 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, + ary: ArrayOrContainerT, + ) -> ArrayOrContainerT: + if not isinstance(ary, DOFArray): + raise TypeError(f"unsupported array container: '{type(ary).__name__}'") + + from pytential.linalg.utils import split_array + targets, sources = hmat.skeletons[0].tgt_src_index + x = split_array(actx.to_numpy(flatten(ary, actx)), sources) + + # NOTE: this computes a telescoping product of the form + # + # A x_0 = (D0 + L0 (D1 + L1 (...) R1) R0) x_0 + # + # with arbitrary numbers of levels. When recursing down, we compute + # + # x_{k + 1} = R_k x_k + # z_{k + 1} = D_k x_k + # + # and, at the root level, we have + # + # x_{N + 1} = z_{N + 1} = D_N x_N. + # + # When recursing back up, we take `b_{N + 1} = x_{N + 1}` and + # + # b_{k - 1} = z_k + L_k b_k + # + # which gives back the desired product when we reach the leaf level again. + + d_dot_x = np.empty(hmat.nlevels, dtype=object) + + # {{{ recurse down + + from pytential.linalg.cluster import cluster + clevels = list(hmat.ctree.levels(root=True)) + + for k, clevel in enumerate(clevels): + skeleton = hmat.skeletons[k] + assert x.shape == (skeleton.nclusters,) + assert skeleton.tgt_src_index.shape[1] == sum([xi.size for xi in x]) + + d_dot_x_k = np.empty(skeleton.nclusters, dtype=object) + r_dot_x_k = np.empty(skeleton.nclusters, dtype=object) + + for i in range(skeleton.nclusters): + r_dot_x_k[i] = skeleton.R[i] @ x[i] + d_dot_x_k[i] = skeleton.D[i] @ x[i] + + d_dot_x[k] = d_dot_x_k + x = cluster(r_dot_x_k, clevel) + + # }}} + + # {{{ root + + # NOTE: at root level, we just multiply with the full diagonal + b = d_dot_x[hmat.nlevels - 1] + assert b.shape == (1,) + + # }}} + + # {{{ recurse up + + from pytential.linalg.cluster import uncluster + + for k, clevel in reversed(list(enumerate(clevels[:-1]))): + skeleton = hmat.skeletons[k] + d_dot_x_k = d_dot_x[k] + assert d_dot_x_k.shape == (skeleton.nclusters,) + + b = uncluster(b, skeleton.skel_tgt_src_index.targets, clevel) + for i in range(skeleton.nclusters): + b[i] = d_dot_x_k[i] + skeleton.L[i] @ b[i] + + assert b.shape == (hmat.nclusters,) + + # }}} + + b = np.concatenate(b)[np.argsort(targets.indices)] + return unflatten(ary, actx.from_numpy(b), actx) + + +def apply_skeleton_inverse_matvec( + actx: PyOpenCLArrayContext, + hmat: ProxyHierarchicalMatrix, + ary: ArrayOrContainerT, + ) -> ArrayOrContainerT: + if not isinstance(ary, DOFArray): + raise TypeError(f"unsupported array container: '{type(ary).__name__}'") + + from pytential.linalg.utils import split_array + targets, sources = hmat.skeletons[0].tgt_src_index + + b = split_array(actx.to_numpy(flatten(ary, actx)), targets) + inv_dhat_dot_b = np.empty(hmat.nlevels, dtype=object) + + # {{{ recurse down + + from pytential.linalg.cluster import cluster + clevels = list(hmat.ctree.levels(root=True)) + + for k, clevel in enumerate(clevels): + skeleton = hmat.skeletons[k] + assert b.shape == (skeleton.nclusters,) + assert skeleton.tgt_src_index.shape[0] == sum([bi.size for bi in b]) + + inv_d_dot_b_k = np.empty(skeleton.nclusters, dtype=object) + inv_dhat_dot_b_k = np.empty(skeleton.nclusters, dtype=object) + + for i in range(skeleton.nclusters): + inv_dhat_dot_b_k[i] = ( + skeleton.Dhat[i] @ (skeleton.R[i] @ (skeleton.invD[i] @ b[i])) + ) + inv_d_dot_b_k[i] = skeleton.invD[i] @ b[i] + + inv_dhat_dot_b[k] = inv_dhat_dot_b_k + b = cluster(inv_dhat_dot_b_k, clevel) + + # }}} + + # {{{ root + + from pytools.obj_array import make_obj_array + assert b.shape == (1,) + x = make_obj_array([ + la.solve(D, bi) for D, bi in zip(hmat.skeletons[-1].D, b) + ]) + + # }}} + + # {{{ recurse up + + from pytential.linalg.cluster import uncluster + + for k, clevel in reversed(list(enumerate(clevels[:-1]))): + skeleton = hmat.skeletons[k] + inv_dhat_dot_b_k0 = inv_dhat_dot_b[k] + inv_dhat_dot_b_k1 = inv_dhat_dot_b[k + 1] + assert inv_d_dot_b_k.shape == (skeleton.nclusters,) + + x = uncluster(x, skeleton.skel_tgt_src_index.sources, clevel) + inv_dhat_dot_b_k1 = uncluster( + inv_dhat_dot_b_k1, skeleton.skel_tgt_src_index.sources, clevel) + + for i in range(skeleton.nclusters): + x[i] = skeleton.invD[i] @ ( + inv_dhat_dot_b_k0[i] + - skeleton.L[i] @ inv_dhat_dot_b_k1[i] + + skeleton.L[i] @ (skeleton.Dhat @ x[i]) + ) + + assert x.shape == (hmat.nclusters,) + + # }}} + + x = np.concatenate(x)[np.argsort(sources.indices)] + return unflatten(ary, actx.from_numpy(x), 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]], *, + method: str = "forward", + 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, + + _approx_nproxy: Optional[int] = None, + _proxy_radius_factor: Optional[float] = None, + ): + if method not in ("forward", "backard"): + raise ValueError(f"unknown matvec method: '{method}'") + + 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.proxy import QBXProxyGenerator + proxy = QBXProxyGenerator(places, + approx_nproxy=_approx_nproxy, + radius_factor=_proxy_radius_factor) + + from pytential.linalg.skeletonization import make_skeletonization_wrangler + wrangler = make_skeletonization_wrangler( + places, exprs, input_exprs, + domains=domains, context=context) + + from pytential.linalg.skeletonization import rec_skeletonize_by_proxy + skeletons = rec_skeletonize_by_proxy( + actx, places, ctree, tgt_src_index, exprs, input_exprs, + id_eps=id_eps, + max_particles_in_box=_max_particles_in_box, + _proxy=proxy, + _wrangler=wrangler, + ) + + if method == "backward": + pass + + return ProxyHierarchicalMatrix( + wrangler=wrangler, ctree=ctree, method=method, skeletons=skeletons) + +# }}} diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py index 2f61b6087..e22eba15c 100644 --- a/pytential/linalg/skeletonization.py +++ b/pytential/linalg/skeletonization.py @@ -26,8 +26,10 @@ TYPE_CHECKING) import numpy as np +import numpy.linalg as la from arraycontext import PyOpenCLArrayContext, Array, ArrayT +from pytools import memoize_method from pytential import GeometryCollection, sym from pytential.linalg.proxy import ProxyGeneratorBase, ProxyClusterGeometryData @@ -637,8 +639,8 @@ def _skeletonize_block_by_proxy_with_mats( ibrow=ibrow, ibcol=ibcol) ) - src_skl_indices = np.empty(nclusters, dtype=object) - tgt_skl_indices = np.empty(nclusters, dtype=object) + src_skel_indices = np.empty(nclusters, dtype=object) + tgt_skel_indices = np.empty(nclusters, dtype=object) skel_starts = np.zeros(nclusters + 1, dtype=np.int32) L = np.empty(nclusters, dtype=object) @@ -661,14 +663,14 @@ def _skeletonize_block_by_proxy_with_mats( assert k > 0 L[i] = interp.T - tgt_skl_indices[i] = tgt_src_index.targets.cluster_indices(i)[idx[:k]] + tgt_skel_indices[i] = tgt_src_index.targets.cluster_indices(i)[idx[:k]] # skeletonize source points k, idx, interp = interp_decomp(src_mat, rank=k, eps=None) assert k > 0 R[i] = interp - src_skl_indices[i] = tgt_src_index.sources.cluster_indices(i)[idx[:k]] + src_skel_indices[i] = tgt_src_index.sources.cluster_indices(i)[idx[:k]] skel_starts[i + 1] = skel_starts[i] + k assert R[i].shape == (k, src_mat.shape[1]) @@ -679,9 +681,9 @@ def _skeletonize_block_by_proxy_with_mats( D = make_flat_cluster_diag(mat, tgt_src_index) from pytential.linalg import TargetAndSourceClusterList, make_index_list - src_skl_index = make_index_list(np.hstack(src_skl_indices), skel_starts) - tgt_skl_index = make_index_list(np.hstack(tgt_skl_indices), skel_starts) - skel_tgt_src_index = TargetAndSourceClusterList(tgt_skl_index, src_skl_index) + src_skel_index = make_index_list(np.hstack(src_skel_indices), skel_starts) + tgt_skel_index = make_index_list(np.hstack(tgt_skel_indices), skel_starts) + skel_tgt_src_index = TargetAndSourceClusterList(tgt_skel_index, src_skel_index) return SkeletonizationResult( L=L, R=R, D=D, @@ -776,6 +778,21 @@ def __post_init__(self): def nclusters(self): return self.tgt_src_index.nclusters + @property + @memoize_method + def invD(self) -> np.ndarray: + from pytools.obj_array import make_obj_array + return make_obj_array([la.inv(D) for D in self.D]) + + @property + @memoize_method + def Dhat(self) -> np.ndarray: + from pytools.obj_array import make_obj_array + return make_obj_array([ + la.inv(self.R[i] @ self.invD[i] @ self.L[i]) + for i in range(self.nclusters) + ]) + def skeletonize_by_proxy( actx: PyOpenCLArrayContext, diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py index d831b04b9..bf0531a39 100644 --- a/pytential/linalg/utils.py +++ b/pytential/linalg/utils.py @@ -88,6 +88,9 @@ def cluster_size(self, i: int) -> int: return self.starts[i + 1] - self.starts[i] + def cluster_slice(self, i: int) -> slice: + return np.s_[self.starts[i]:self.starts[i + 1]] + def cluster_indices(self, i: int) -> np.ndarray: """ :returns: a view into the :attr:`indices` array for the range @@ -97,7 +100,7 @@ def cluster_indices(self, i: int) -> np.ndarray: raise IndexError( f"cluster {i} is out of bounds for {self.nclusters} clusters") - return self.indices[self.starts[i]:self.starts[i + 1]] + return self.indices[self.cluster_slice(i)] def cluster_take(self, x: np.ndarray, i: int) -> np.ndarray: """ @@ -334,6 +337,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) + ]) + # }}} @@ -440,6 +452,22 @@ def mnorm(x: np.ndarray, y: np.ndarray) -> float: return tgt_error, src_error +def skeletonization_matrix( + mat: np.ndarray, skeleton: "SkeletonizationResult", + ) -> Tuple[np.ndarray, np.ndarray]: + D = np.empty(skeleton.nclusters, dtype=object) + S = np.empty((skeleton.nclusters, skeleton.nclusters), dtype=object) + + from itertools import product + for i, j in product(range(skeleton.nclusters), repeat=2): + if i == j: + D[i] = skeleton.tgt_src_index.cluster_take(mat, i, i) + else: + S[i, j] = skeleton.skel_tgt_src_index.cluster_take(mat, i, j) + + return D, S + + def skeletonization_error( mat: np.ndarray, skeleton: "SkeletonizationResult", *, ord: Optional[float] = None, diff --git a/test/test_linalg_hmatrix.py b/test/test_linalg_hmatrix.py new file mode 100644 index 000000000..3e7cad567 --- /dev/null +++ b/test/test_linalg_hmatrix.py @@ -0,0 +1,147 @@ +__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 replace +import pytest + +from pytential import bind, sym +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, + ]) + + +HMATRIX_TEST_CASES = [ + extra.CurveTestCase( + name="starfish", + op_type="scalar", + target_order=4, + curve_fn=NArmedStarfish(5, 0.25), + resolutions=[32]), + extra.CurveTestCase( + name="starfish", + op_type="double", + 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) +@pytest.mark.parametrize("discr_stage", [ + sym.QBX_SOURCE_STAGE1, sym.QBX_SOURCE_STAGE2 + ]) +def test_hmatrix_forward_matvec(actx_factory, case, discr_stage, visualize=False): + actx = actx_factory() + + if visualize: + logging.basicConfig(level=logging.INFO) + + case = replace(case, + tree_kind=None, id_eps=1.0e-12, skel_discr_stage=discr_stage) + 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("dd %s", 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_hmap = hmat @ x + r_lpot = bind(places, sym_op, auto_where=dd)(actx, u=x) + + from meshmode.dof_array import flat_norm + error = actx.to_numpy( + flat_norm(r_hmap - r_lpot) / flat_norm(r_lpot) + ) + logger.info("error: %.12e", error) + assert error < 0.1 + + # }}} + + if not visualize: + return + + if case.ambient_dim == 2: + import matplotlib.pyplot as mp + fig = mp.figure(figsize=(10, 10), dpi=300) + ax = fig.gca() + + from arraycontext import flatten + r_hmap = actx.to_numpy(flatten(r_hmap, actx)) + r_lpot = actx.to_numpy(flatten(r_lpot, actx)) + + ax.plot(r_hmap) + ax.plot(r_lpot, "k--") + + fig.savefig("test_linalg_hmatrix") + + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: fdm=marker