implement hierachical matrix solver

Author: alexfikl

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
+
 # }}}
 
 

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)
+
+# }}}