diff --git a/.gitignore b/.gitignore
index daf719e6f..c89b827fb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -21,6 +21,8 @@ examples/*.pdf
*.vtu
*.vts
+*.png
+*.gif
.cache
diff --git a/examples/scaling-study-hmatrix.py b/examples/scaling-study-hmatrix.py
new file mode 100644
index 000000000..d163e7383
--- /dev/null
+++ b/examples/scaling-study-hmatrix.py
@@ -0,0 +1,198 @@
+__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
+
+import numpy as np
+
+from meshmode.array_context import PyOpenCLArrayContext
+from pytools.convergence import EOCRecorder
+
+from pytential import GeometryCollection, sym
+
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class Timings:
+ build: float
+ matvec: float
+
+
+def run_hmatrix_matvec(
+ actx: PyOpenCLArrayContext,
+ places: GeometryCollection, *,
+ dofdesc: sym.DOFDescriptor) -> None:
+ from sumpy.kernel import LaplaceKernel
+ kernel = LaplaceKernel(places.ambient_dim)
+ sym_u = sym.var("u")
+ sym_op = -0.5 * sym_u + sym.D(kernel, sym_u, qbx_forced_limit="avg")
+
+ density_discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage)
+ u = actx.thaw(density_discr.nodes()[0])
+
+ def build_hmat():
+ from pytential.linalg.hmatrix import build_hmatrix_by_proxy
+ return build_hmatrix_by_proxy(
+ actx, places, sym_op, sym_u,
+ domains=[dofdesc],
+ context={},
+ auto_where=dofdesc,
+ id_eps=1.0e-10,
+ _tree_kind="adaptive-level-restricted",
+ _approx_nproxy=64,
+ _proxy_radius_factor=1.15).get_forward()
+
+ # warmup
+ from pytools import ProcessTimer
+ with ProcessTimer() as pt:
+ hmat = build_hmat()
+ actx.queue.finish()
+
+ logger.info("build(warmup): %s", pt)
+
+ # build
+ with ProcessTimer() as pt:
+ hmat = build_hmat()
+ actx.queue.finish()
+
+ t_build = pt.wall_elapsed
+ logger.info("build: %s", pt)
+
+ # matvec
+ with ProcessTimer() as pt:
+ du = hmat @ u
+ assert du is not None
+ actx.queue.finish()
+
+ t_matvec = pt.wall_elapsed
+ logger.info("matvec: %s", pt)
+
+ return Timings(t_build, t_matvec)
+
+
+def run_scaling_study(
+ ambient_dim: int, *,
+ target_order: int = 4,
+ source_ovsmp: int = 4,
+ qbx_order: int = 4,
+ ) -> None:
+ dd = sym.DOFDescriptor(f"d{ambient_dim}", discr_stage=sym.QBX_SOURCE_STAGE2)
+
+ import pyopencl as cl
+ ctx = cl.create_some_context()
+ queue = cl.CommandQueue(ctx)
+ actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
+
+ eoc_build = EOCRecorder()
+ eoc_matvec = EOCRecorder()
+
+ import meshmode.mesh.generation as mgen
+ import meshmode.discretization.poly_element as mpoly
+ resolutions = [64, 128, 256, 512, 1024, 1536, 2048, 2560, 3072]
+
+ for n in resolutions:
+ mesh = mgen.make_curve_mesh(
+ mgen.NArmedStarfish(5, 0.25),
+ np.linspace(0, 1, n),
+ order=target_order)
+
+ from meshmode.discretization import Discretization
+ pre_density_discr = Discretization(actx, mesh,
+ mpoly.InterpolatoryQuadratureGroupFactory(target_order))
+
+ from pytential.qbx import QBXLayerPotentialSource
+ qbx = QBXLayerPotentialSource(
+ pre_density_discr,
+ fine_order=source_ovsmp * target_order,
+ qbx_order=qbx_order,
+ fmm_order=False, fmm_backend=None,
+ )
+ places = GeometryCollection(qbx, auto_where=dd.geometry)
+ density_discr = places.get_discretization(dd.geometry, dd.discr_stage)
+
+ logger.info("ndofs: %d", density_discr.ndofs)
+ logger.info("nelements: %d", density_discr.mesh.nelements)
+
+ timings = run_hmatrix_matvec(actx, places, dofdesc=dd)
+ eoc_build.add_data_point(density_discr.ndofs, timings.build)
+ eoc_matvec.add_data_point(density_discr.ndofs, timings.matvec)
+
+ for name, eoc in [("build", eoc_build), ("matvec", eoc_matvec)]:
+ logger.info("%s\n%s",
+ name, eoc.pretty_print(
+ abscissa_label="dofs",
+ error_label=f"{name} (s)",
+ abscissa_format="%d",
+ error_format="%.3fs",
+ eoc_format="%.2f",
+ )
+ )
+ visualize_eoc(f"scaling-study-hmatrix-{name}", eoc, 1)
+
+
+def visualize_eoc(
+ filename: str, eoc: EOCRecorder, order: int,
+ overwrite: bool = False) -> None:
+ try:
+ import matplotlib.pyplot as plt
+ except ImportError:
+ logger.info("matplotlib not available for plotting")
+ return
+
+ fig = plt.figure(figsize=(10, 10), dpi=300)
+ ax = fig.gca()
+
+ h, error = np.array(eoc.history).T # type: ignore[no-untyped-call]
+ ax.loglog(h, error, "o-")
+
+ max_h = np.max(h)
+ min_e = np.min(error)
+ max_e = np.max(error)
+ min_h = np.exp(np.log(max_h) + np.log(min_e / max_e) / order)
+
+ ax.loglog(
+ [max_h, min_h], [max_e, min_e], "k-", label=rf"$\mathcal{{O}}(h^{order})$"
+ )
+
+ # }}}
+
+ ax.grid(True, which="major", linestyle="-", alpha=0.75)
+ ax.grid(True, which="minor", linestyle="--", alpha=0.5)
+
+ ax.set_xlabel("$N$")
+ ax.set_ylabel("$T~(s)$")
+
+ import pathlib
+ filename = pathlib.Path(filename)
+ if not overwrite and filename.exists():
+ raise FileExistsError(f"output file '{filename}' already exists")
+
+ fig.savefig(filename)
+ plt.close(fig)
+
+
+if __name__ == "__main__":
+ run_scaling_study(ambient_dim=2)
diff --git a/pytential/linalg/direct_solver_symbolic.py b/pytential/linalg/direct_solver_symbolic.py
index c57752840..5c11eee1d 100644
--- a/pytential/linalg/direct_solver_symbolic.py
+++ b/pytential/linalg/direct_solver_symbolic.py
@@ -54,12 +54,16 @@ def prepare_expr(places, exprs, auto_where=None):
return _prepare_expr(places, exprs, auto_where=auto_where)
-def prepare_proxy_expr(places, exprs, auto_where=None):
+def prepare_proxy_expr(
+ places, exprs,
+ auto_where=None,
+ remove_transforms: bool = True):
def _prepare_expr(expr):
# remove all diagonal / non-operator terms in the expression
expr = IntGTermCollector()(expr)
# ensure all IntGs remove all the kernel derivatives
- expr = KernelTransformationRemover()(expr)
+ if remove_transforms:
+ expr = KernelTransformationRemover()(expr)
# ensure all IntGs have their source and targets set
expr = DOFDescriptorReplacer(auto_where[0], auto_where[1])(expr)
diff --git a/pytential/linalg/hmatrix.py b/pytential/linalg/hmatrix.py
new file mode 100644
index 000000000..a72a209cb
--- /dev/null
+++ b/pytential/linalg/hmatrix.py
@@ -0,0 +1,544 @@
+__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, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import numpy.linalg as la
+from scipy.sparse.linalg import LinearOperator
+
+from arraycontext import PyOpenCLArrayContext, ArrayOrContainerT, flatten, unflatten
+from meshmode.dof_array import DOFArray
+
+from pytential import GeometryCollection, sym
+from pytential.linalg.proxy import ProxyGeneratorBase
+from pytential.linalg.cluster import ClusterTree, ClusterLevel
+from pytential.linalg.skeletonization import (
+ SkeletonizationWrangler, SkeletonizationResult)
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+__doc__ = """
+Hierarical Matrix Construction
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ProxyHierarchicalMatrixWrangler
+.. autoclass:: ProxyHierarchicalMatrix
+.. autoclass:: ProxyHierarchicalForwardMatrix
+.. autoclass:: ProxyHierarchicalBackwardMatrix
+
+.. autofunction:: build_hmatrix_by_proxy
+"""
+
+
+# {{{ error model
+
+def hmatrix_error_from_param(
+ ambient_dim: int,
+ *,
+ id_eps: float,
+ id_rank: int,
+ min_proxy_radius: float,
+ max_cluster_radius: float,
+ nproxies: int,
+ nsources: int,
+ ntargets: int, c: float = 1.0e-3) -> float:
+ if ambient_dim == 2:
+ p = int(0.5 * id_rank)
+ elif ambient_dim == 3:
+ p = int((np.sqrt(1 + 4 * id_rank) - 1) / 2)
+ else:
+ raise ValueError(f"unsupported ambient dimension: '{ambient_dim}'")
+
+ rho = alpha = max_cluster_radius / min_proxy_radius
+ return (
+ c * rho ** (p + 1) / (1 - rho)
+ + np.sqrt(nsources / nproxies)
+ * (1 - alpha ** (p + 1)) / (1 - alpha) * id_eps
+ )
+
+# }}}
+
+
+# {{{ update diagonals
+
+def _update_skeleton_diagonal(
+ skeleton: SkeletonizationResult,
+ parent: Optional[SkeletonizationResult],
+ clevel: Optional[ClusterLevel],
+ diagonal: Optional[np.ndarray] = None) -> SkeletonizationResult:
+ """Due to the evaluation in :func:`_skeletonize_block_by_proxy_with_mats`,
+ the diagonal matrix in *skeleton* also contains the indices from its
+ parent. In particular, at a level :math:`l` we need the diagonal block::
+
+ 0 D_{i, j + 1} D_{i, j + 2}
+ D_{i + 1, j} 0 D_{i + 1, j + 2}
+ D_{i + 2, j} D_{i + 2, j + 1} 0
+
+ but the version in *skeleton* also fills in the 0 blocks in there. This
+ routine goes through them and zeros them out.
+ """
+
+ if clevel is None:
+ return skeleton
+
+ assert parent is not None
+ assert skeleton.tgt_src_index.shape == parent.skel_tgt_src_index.shape
+
+ if diagonal is None:
+ diagonal = np.zeros(parent.nclusters)
+
+ from numbers import Number
+ if isinstance(diagonal, Number):
+ diagonal = np.full(parent.nclusters, diagonal)
+
+ assert diagonal.size == parent.nclusters
+ targets, sources = parent.skel_tgt_src_index
+
+ # FIXME: nicer way to do this?
+ mat: np.ndarray = np.empty(skeleton.nclusters, dtype=object)
+ for k in range(skeleton.nclusters):
+ D = skeleton.D[k].copy()
+
+ i = j = 0
+ for icluster in clevel.parent_map[k]:
+ di = targets.cluster_size(icluster)
+ dj = sources.cluster_size(icluster)
+ D[np.s_[i:i + di], np.s_[j:j + dj]] = diagonal[icluster]
+
+ i += di
+ j += dj
+
+ assert D.shape == (i, j)
+ mat[k] = D
+
+ from dataclasses import replace
+ return replace(skeleton, D=mat)
+
+
+def _update_skeletons_diagonal(
+ wrangler: "ProxyHierarchicalMatrixWrangler",
+ func: Callable[[SkeletonizationResult], Optional[np.ndarray]],
+ ) -> np.ndarray:
+ skeletons: np.ndarray = np.empty(wrangler.skeletons.shape, dtype=object)
+ skeletons[0] = wrangler.skeletons[0]
+
+ for i in range(1, wrangler.ctree.nlevels):
+ skeletons[i] = _update_skeleton_diagonal(
+ wrangler.skeletons[i],
+ wrangler.skeletons[i - 1],
+ wrangler.ctree.levels[i - 1],
+ diagonal=func(skeletons[i - 1]))
+
+ return skeletons
+
+# }}}
+
+
+# {{{ ProxyHierarchicalMatrix
+
+@dataclass(frozen=True)
+class ProxyHierarchicalMatrixWrangler:
+ """
+ .. automethod:: get_forward
+ .. automethod:: get_backward
+ """
+
+ wrangler: SkeletonizationWrangler
+ proxy: ProxyGeneratorBase
+ ctree: ClusterTree
+ skeletons: np.ndarray
+
+ def get_forward(self) -> "ProxyHierarchicalForwardMatrix":
+ return ProxyHierarchicalForwardMatrix(
+ ctree=self.ctree,
+ skeletons=_update_skeletons_diagonal(self, lambda s: None),
+ )
+
+ def get_backward(self) -> "ProxyHierarchicalBackwardMatrix":
+ return ProxyHierarchicalBackwardMatrix(
+ ctree=self.ctree,
+ skeletons=_update_skeletons_diagonal(self, lambda s: s.Dhat))
+
+
+@dataclass(frozen=True)
+class ProxyHierarchicalMatrix(LinearOperator):
+ """
+ .. 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
+
+ @property
+ def nclusters(self):
+ return self.skeletons[0].nclusters
+
+ def __matmul__(self, x: ArrayOrContainerT) -> ArrayOrContainerT:
+ """Same as :meth:`matvec`."""
+ return self._matvec(x)
+
+ def _matmat(self, mat):
+ raise NotImplementedError
+
+ def _adjoint(self, x):
+ raise NotImplementedError
+
+# }}}
+
+
+# {{{ forward
+
+@dataclass(frozen=True)
+class ProxyHierarchicalForwardMatrix(ProxyHierarchicalMatrix):
+ def _matvec(self, x: ArrayOrContainerT) -> ArrayOrContainerT:
+ if isinstance(x, DOFArray):
+ 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")
+
+ ary = actx.to_numpy(flatten(x, actx))
+ elif isinstance(x, np.ndarray) and x.dtype.char != "O":
+ ary = x
+ else:
+ raise TypeError(f"unsupported input type: {type(x)}")
+
+ assert actx is None or isinstance(actx, PyOpenCLArrayContext)
+ result = apply_skeleton_forward_matvec(self, ary)
+
+ if isinstance(x, DOFArray):
+ assert actx is not None
+ result = unflatten(x, actx.from_numpy(result), actx)
+
+ return result # type: ignore[return-value]
+
+
+def apply_skeleton_forward_matvec(
+ hmat: ProxyHierarchicalMatrix,
+ ary: ArrayOrContainerT,
+ ) -> ArrayOrContainerT:
+ from pytential.linalg.cluster import split_array
+ targets, sources = hmat.skeletons[0].tgt_src_index
+ x = split_array(ary, sources) # type: ignore[arg-type]
+
+ # 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.ndarray = np.empty(hmat.nlevels, dtype=object)
+
+ # {{{ recurse down
+
+ from pytential.linalg.cluster import cluster
+ for k, clevel in enumerate(hmat.ctree.levels):
+ 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.ndarray = np.empty(skeleton.nclusters, dtype=object)
+ r_dot_x_k: np.ndarray = 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(hmat.ctree.levels[:-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,)
+
+ # }}}
+
+ return np.concatenate(b)[np.argsort(targets.indices)]
+
+# }}}
+
+
+# {{{ backward
+
+@dataclass(frozen=True)
+class ProxyHierarchicalBackwardMatrix(ProxyHierarchicalMatrix):
+ def _matvec(self, x: ArrayOrContainerT) -> ArrayOrContainerT:
+ if isinstance(x, DOFArray):
+ 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")
+
+ ary = actx.to_numpy(flatten(x, actx))
+ elif isinstance(x, np.ndarray) and x.dtype.char != "O":
+ ary = x
+ else:
+ raise TypeError(f"unsupported input type: {type(x)}")
+
+ assert actx is None or isinstance(actx, PyOpenCLArrayContext)
+ result = apply_skeleton_backward_matvec(actx, self, ary)
+
+ if isinstance(x, DOFArray):
+ assert actx is not None
+ result = unflatten(x, actx.from_numpy(result), actx)
+
+ return result # type: ignore[return-value]
+
+
+def apply_skeleton_backward_matvec(
+ actx: Optional[PyOpenCLArrayContext],
+ hmat: ProxyHierarchicalMatrix,
+ ary: ArrayOrContainerT,
+ ) -> ArrayOrContainerT:
+ from pytential.linalg.cluster import split_array
+ targets, sources = hmat.skeletons[0].tgt_src_index
+
+ b = split_array(ary, targets) # type: ignore[arg-type]
+ r_dot_b: np.ndarray = np.empty(hmat.nlevels, dtype=object)
+
+ # {{{ recurse down
+
+ # NOTE: this solves a telescoping product of the form
+ #
+ # A x_0 = (D0 + L0 (D1 + L1 (...) R1) R0) x_0 = b_0
+ #
+ # with arbitrary numbers of levels. When recursing down, we compute
+ #
+ # b_{k + 1} = \hat{D}_k R_k D_k^{-1} b_k
+ # \hat{D}_k = (R_k D_k^{-1} L_k)^{-1}
+ #
+ # and, at the root level, we solve
+ #
+ # D_N x_N = b_N.
+ #
+ # When recursing back up, we take `b_{N + 1} = x_{N + 1}` and
+ #
+ # x_{k} = D_k^{-1} (b_k - L_k b_{k + 1} + L_k \hat{D}_k x_{k + 1})
+ #
+ # which gives back the desired product when we reach the leaf level again.
+
+ from pytential.linalg.cluster import cluster
+
+ for k, clevel in enumerate(hmat.ctree.levels):
+ skeleton = hmat.skeletons[k]
+ assert b.shape == (skeleton.nclusters,)
+ assert skeleton.tgt_src_index.shape[0] == sum(bi.size for bi in b)
+
+ dhat_dot_b_k: np.ndarray = np.empty(skeleton.nclusters, dtype=object)
+ for i in range(skeleton.nclusters):
+ dhat_dot_b_k[i] = (
+ skeleton.Dhat[i] @ (skeleton.R[i] @ (skeleton.invD[i] @ b[i]))
+ )
+
+ r_dot_b[k] = b
+ b = cluster(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(hmat.ctree.levels[:-1]))):
+ skeleton = hmat.skeletons[k]
+ b0 = r_dot_b[k]
+ b1 = r_dot_b[k + 1]
+ assert b0.shape == (skeleton.nclusters,)
+
+ x = uncluster(x, skeleton.skel_tgt_src_index.sources, clevel)
+ b1 = uncluster(b1, skeleton.skel_tgt_src_index.targets, clevel)
+
+ for i in range(skeleton.nclusters):
+ sx = b1[i] - skeleton.Dhat[i] @ x[i]
+ x[i] = skeleton.invD[i] @ (b0[i] - skeleton.L[i] @ sx)
+
+ assert x.shape == (hmat.nclusters,)
+
+ # }}}
+
+ return np.concatenate(x)[np.argsort(sources.indices)]
+
+# }}}
+
+
+# {{{ build_hmatrix_by_proxy
+
+def build_hmatrix_by_proxy(
+ actx: PyOpenCLArrayContext,
+ places: GeometryCollection,
+ exprs: Union[sym.Expression, Iterable[sym.Expression]],
+ input_exprs: Union[sym.Expression, Iterable[sym.Expression]], *,
+ auto_where: Optional[sym.DOFDescriptorLike] = None,
+ domains: Optional[Sequence[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!
+ _tree_kind: Optional[str] = "adaptive-level-restricted",
+ _weighted_proxy: Optional[Union[bool, Tuple[bool, bool]]] = None,
+
+ # 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?
+ _max_particles_in_box: Optional[int] = None,
+ _approx_nproxy: Optional[int] = None,
+ _proxy_radius_factor: Optional[float] = None,
+ ) -> ProxyHierarchicalMatrixWrangler:
+ from pytential.symbolic.matrix import P2PClusterMatrixBuilder
+ from pytential.linalg.skeletonization import make_skeletonization_wrangler
+
+ def P2PClusterMatrixBuilderWithDiagonal(*args, **kwargs):
+ kwargs["exclude_self"] = True
+ return P2PClusterMatrixBuilder(*args, **kwargs)
+
+ wrangler = make_skeletonization_wrangler(
+ places, exprs, input_exprs,
+ domains=domains, context=context, auto_where=auto_where,
+ _weighted_proxy=_weighted_proxy,
+ # _neighbor_cluster_builder=P2PClusterMatrixBuilderWithDiagonal,
+ # _proxy_source_cluster_builder=P2PClusterMatrixBuilder,
+ # _proxy_target_cluster_builder=P2PClusterMatrixBuilder,
+ )
+
+ if wrangler.nrows != 1 or wrangler.ncols != 1:
+ raise ValueError("multi-block operators are not supported")
+
+ from pytential.linalg.proxy import QBXProxyGenerator
+ proxy = QBXProxyGenerator(places,
+ approx_nproxy=_approx_nproxy,
+ radius_factor=_proxy_radius_factor)
+
+ from pytential.linalg.cluster import partition_by_nodes
+ cluster_index, ctree = partition_by_nodes(
+ actx, places,
+ dofdesc=wrangler.domains[0],
+ tree_kind=_tree_kind,
+ max_particles_in_box=_max_particles_in_box)
+
+ logger.info("tree levels: %d", ctree.nlevels)
+ logger.info("cluster count: %d", cluster_index.nclusters)
+ logger.info("root cluster sizes: %s", [
+ # NOTE: making into a list so that they all get printed
+ int(s) for s in np.diff(cluster_index.starts)
+ ])
+
+ 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,
+ id_eps=id_eps,
+ max_particles_in_box=_max_particles_in_box,
+ _proxy=proxy,
+ _wrangler=wrangler,
+ )
+
+ return ProxyHierarchicalMatrixWrangler(
+ wrangler=wrangler, proxy=proxy, ctree=ctree, skeletons=skeletons
+ )
+
+# }}}
+
+# vim: foldmethod=marker
diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py
index 5e3b9cec4..6eb8c721f 100644
--- a/pytential/linalg/proxy.py
+++ b/pytential/linalg/proxy.py
@@ -360,7 +360,7 @@ def __call__(self,
source_dd.geometry, source_dd.discr_stage)
assert isinstance(discr, Discretization)
- include_cluster_radii = kwargs.pop("include_cluster_radii", False)
+ include_cluster_radii = kwargs.pop("include_cluster_radii", True)
# {{{ get proxy centers and radii
diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py
index e9ebec2eb..477acf9c8 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
+from pytools import memoize_method
from pytential import GeometryCollection, sym
from pytential.linalg.proxy import ProxyGeneratorBase, ProxyClusterGeometryData
@@ -425,6 +427,7 @@ def make_skeletonization_wrangler(
# internal
_weighted_proxy: Optional[Union[bool, Tuple[bool, bool]]] = None,
+ _remove_source_transforms: bool = False,
_proxy_source_cluster_builder: Optional[Callable[..., np.ndarray]] = None,
_proxy_target_cluster_builder: Optional[Callable[..., np.ndarray]] = None,
_neighbor_cluster_builder: Optional[Callable[..., np.ndarray]] = None,
@@ -451,9 +454,13 @@ def make_skeletonization_wrangler(
exprs = prepare_expr(places, exprs, auto_where)
source_proxy_exprs = prepare_proxy_expr(
- places, exprs, (auto_where[0], PROXY_SKELETONIZATION_TARGET))
+ places, exprs, (auto_where[0], PROXY_SKELETONIZATION_TARGET),
+ remove_transforms=_remove_source_transforms)
target_proxy_exprs = prepare_proxy_expr(
- places, exprs, (PROXY_SKELETONIZATION_SOURCE, auto_where[1]))
+ places, exprs, (PROXY_SKELETONIZATION_SOURCE, auto_where[1]),
+ # NOTE: transforms are unconditionally removed here because the
+ # source would be the proxies, where we do not have normals, etc.
+ remove_transforms=True)
# }}}
@@ -463,7 +470,7 @@ def make_skeletonization_wrangler(
weighted_sources = weighted_targets = True
elif isinstance(_weighted_proxy, bool):
weighted_sources = weighted_targets = _weighted_proxy
- elif isinstance(_weighted_proxy, tuple):
+ elif isinstance(_weighted_proxy, tuple) and len(_weighted_proxy) == 2:
weighted_sources, weighted_targets = _weighted_proxy
else:
raise ValueError(f"unknown value for weighting: '{_weighted_proxy}'")
@@ -661,9 +668,9 @@ def _skeletonize_block_by_proxy_with_mats(
if __debug__:
isfinite = np.isfinite(tgt_mat)
- assert np.all(isfinite), np.where(isfinite)
+ assert np.all(isfinite), np.where(~isfinite)
isfinite = np.isfinite(src_mat)
- assert np.all(isfinite), np.where(isfinite)
+ assert np.all(isfinite), np.where(~isfinite)
# skeletonize target points
k, idx, interp = interp_decomp(tgt_mat.T, rank=k, eps=id_eps)
@@ -813,6 +820,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 1a8da49ff..4c6f19d40 100644
--- a/pytential/linalg/utils.py
+++ b/pytential/linalg/utils.py
@@ -440,6 +440,22 @@ def mnorm(x: np.ndarray, y: np.ndarray) -> "np.floating[Any]":
return tgt_error, src_error
+def skeletonization_matrix(
+ mat: np.ndarray, skeleton: "SkeletonizationResult",
+ ) -> Tuple[np.ndarray, np.ndarray]:
+ D: np.ndarray = np.empty(skeleton.nclusters, dtype=object)
+ S: np.ndarray = 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,
@@ -501,4 +517,50 @@ def skeletonization_error(
# }}}
+
+# {{{ eigenvalues
+
+def eigs(
+ mat, *,
+ k: int = 6,
+ which: str = "LM",
+ maxiter: Optional[int] = None,
+ tol: float = 0.0) -> np.ndarray:
+ import scipy.sparse.linalg as ssla
+
+ result = ssla.eigs(mat,
+ k=k,
+ which=which,
+ maxiter=maxiter,
+ tol=tol,
+ return_eigenvectors=False)
+
+ imag_norm = np.linalg.norm(np.imag(result), ord=np.inf)
+ if imag_norm > 1.0e-14:
+ from warnings import warn
+ warn(f"eigenvalues are not real enough: norm(imag) = {imag_norm:.12e}")
+
+ return result
+
+
+def cond(mat, *,
+ mat_inv=None,
+ p: Optional[float] = None,
+ tol: float = 1.0e-6) -> float:
+ if p is None:
+ p = 2
+
+ if p != 2:
+ raise ValueError(f"unsupported norm order: '{p}'")
+
+ lambda_max = eigs(mat, k=1, which="LM", tol=tol)
+ if mat_inv is None:
+ lambda_min = eigs(mat, k=1, which="SM", tol=tol)
+ else:
+ lambda_min = eigs(mat_inv, k=1, which="LM", tol=tol)
+
+ return np.abs(lambda_max) / np.abs(lambda_min)
+
+# }}}
+
# vim: foldmethod=marker
diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py
index c24353dca..0484527bb 100644
--- a/pytential/symbolic/matrix.py
+++ b/pytential/symbolic/matrix.py
@@ -461,7 +461,6 @@ def map_int_g(self, expr):
expr.target.geometry, expr.target.discr_stage)
actx = self.array_context
- target_base_kernel = expr.target_kernel.get_base_kernel()
result = 0
for density, kernel in zip(expr.densities, expr.source_kernels):
@@ -475,12 +474,10 @@ def map_int_g(self, expr):
# {{{ generator
- base_kernel = kernel.get_base_kernel()
-
from sumpy.p2p import P2PMatrixGenerator
mat_gen = P2PMatrixGenerator(actx.context,
- source_kernels=(base_kernel,),
- target_kernels=(target_base_kernel,),
+ source_kernels=(kernel,),
+ target_kernels=(expr.target_kernel,),
exclude_self=self.exclude_self)
# }}}
@@ -490,7 +487,7 @@ def map_int_g(self, expr):
# {{{ kernel args
# NOTE: copied from pytential.symbolic.primitives.IntG
- kernel_args = base_kernel.get_args() + base_kernel.get_source_args()
+ kernel_args = kernel.get_args() + kernel.get_source_args()
kernel_args = {arg.loopy_arg.name for arg in kernel_args}
kernel_args = _get_layer_potential_args(
@@ -638,7 +635,6 @@ def map_int_g(self, expr):
expr.target.geometry, expr.target.discr_stage)
actx = self.array_context
- target_base_kernel = expr.target_kernel.get_base_kernel()
result = 0
for kernel, density in zip(expr.source_kernels, expr.densities):
@@ -659,12 +655,10 @@ def map_int_g(self, expr):
# {{{ generator
- base_kernel = kernel.get_base_kernel()
-
from sumpy.p2p import P2PMatrixSubsetGenerator
mat_gen = P2PMatrixSubsetGenerator(actx.context,
- source_kernels=(base_kernel,),
- target_kernels=(target_base_kernel,),
+ source_kernels=(kernel,),
+ target_kernels=(expr.target_kernel,),
exclude_self=self.exclude_self)
# }}}
@@ -674,7 +668,7 @@ def map_int_g(self, expr):
# {{{ kernel args
# NOTE: copied from pytential.symbolic.primitives.IntG
- kernel_args = base_kernel.get_args() + base_kernel.get_source_args()
+ kernel_args = kernel.get_args() + kernel.get_source_args()
kernel_args = {arg.loopy_arg.name for arg in kernel_args}
kernel_args = _get_layer_potential_args(
diff --git a/test/extra_matrix_data.py b/test/extra_matrix_data.py
index 70710bec6..a0018ccce 100644
--- a/test/extra_matrix_data.py
+++ b/test/extra_matrix_data.py
@@ -43,14 +43,23 @@ class MatrixTestCaseMixin:
proxy_target_cluster_builder: Optional[Callable[..., Any]] = None
neighbor_cluster_builder: Optional[Callable[..., Any]] = None
- def get_cluster_index(self, actx, places, dofdesc=None):
+ def max_particles_in_box_for_discr(self, discr):
+ max_particles_in_box = self.max_particles_in_box
+ if max_particles_in_box is None:
+ max_particles_in_box = discr.ndofs // self.approx_cluster_count
+
+ return max_particles_in_box
+
+ def get_cluster_index(
+ self, actx, places, dofdesc=None, max_particles_in_box=None):
if dofdesc is None:
dofdesc = places.auto_source
discr = places.get_discretization(dofdesc.geometry)
- max_particles_in_box = self.max_particles_in_box
if max_particles_in_box is None:
- max_particles_in_box = discr.ndofs // self.approx_cluster_count
+ max_particles_in_box = self.max_particles_in_box
+ if max_particles_in_box is None:
+ max_particles_in_box = discr.ndofs // self.approx_cluster_count
from pytential.linalg.cluster import partition_by_nodes
cindex, ctree = partition_by_nodes(actx, places,
@@ -77,9 +86,11 @@ def get_cluster_index(self, actx, places, dofdesc=None):
return cindex, ctree
- def get_tgt_src_cluster_index(self, actx, places, dofdesc=None):
+ def get_tgt_src_cluster_index(
+ self, actx, places, dofdesc=None, max_particles_in_box=None):
from pytential.linalg import TargetAndSourceClusterList
- cindex, ctree = self.get_cluster_index(actx, places, dofdesc=dofdesc)
+ cindex, ctree = self.get_cluster_index(
+ actx, places, dofdesc=dofdesc, max_particles_in_box=max_particles_in_box)
return TargetAndSourceClusterList(cindex, cindex), ctree
def get_operator(self, ambient_dim, qbx_forced_limit=_NoArgSentinel):
diff --git a/test/test_linalg_hmatrix.py b/test/test_linalg_hmatrix.py
new file mode 100644
index 000000000..1bba7b9d2
--- /dev/null
+++ b/test/test_linalg_hmatrix.py
@@ -0,0 +1,582 @@
+__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
+
+import numpy as np
+
+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=[512]),
+ extra.CurveTestCase(
+ name="starfish",
+ op_type="double",
+ target_order=4,
+ curve_fn=NArmedStarfish(5, 0.25),
+ resolutions=[512]),
+ extra.TorusTestCase(
+ target_order=4,
+ op_type="scalar",
+ resolutions=[0])
+ ]
+
+
+# {{{ test_hmatrix_forward_matvec_single_level
+
+def hmatrix_matvec_single_level(mat, x, skeleton):
+ from pytential.linalg.cluster import split_array
+ targets, sources = skeleton.tgt_src_index
+ y = split_array(x, sources)
+
+ y_hat = np.empty(y.shape, dtype=object)
+
+ for i in range(skeleton.nclusters):
+ y_hat[i] = skeleton.R[i] @ y[i]
+
+ from pytential.linalg.utils import skeletonization_matrix
+ D, S = skeletonization_matrix(mat, skeleton)
+ syhat = np.zeros(y.shape, dtype=object)
+
+ from itertools import product
+ for i, j in product(range(skeleton.nclusters), repeat=2):
+ if i == j:
+ continue
+
+ syhat[i] = syhat[i] + S[i, j] @ y_hat[j]
+
+ for i in range(skeleton.nclusters):
+ y[i] = D[i] @ y[i] + skeleton.L[i] @ syhat[i]
+
+ return np.concatenate(y)[np.argsort(targets.indices)]
+
+
+@pytest.mark.parametrize("case", HMATRIX_TEST_CASES)
+@pytest.mark.parametrize("discr_stage", [sym.QBX_SOURCE_STAGE1])
+def test_hmatrix_forward_matvec_single_level(
+ actx_factory, case, discr_stage, visualize=False):
+ actx = actx_factory()
+
+ if visualize:
+ logging.basicConfig(level=logging.INFO)
+
+ if case.ambient_dim == 2:
+ kwargs = {"proxy_approx_count": 64, "proxy_radius_factor": 1.15}
+ else:
+ kwargs = {"proxy_approx_count": 256, "proxy_radius_factor": 1.25}
+
+ case = replace(case, skel_discr_stage=discr_stage, **kwargs)
+ 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[-1], 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 reference
+
+ from pytential.linalg.direct_solver_symbolic import prepare_expr
+ from pytential.symbolic.matrix import MatrixBuilder
+ sym_u, sym_op = case.get_operator(places.ambient_dim)
+ mat = MatrixBuilder(
+ actx,
+ dep_expr=sym_u,
+ other_dep_exprs=[],
+ dep_discr=density_discr,
+ places=places,
+ context={},
+ )(prepare_expr(places, sym_op, (dd, dd)))
+
+ from arraycontext import flatten, unflatten
+ x = actx.thaw(density_discr.nodes()[0])
+ y = actx.to_numpy(flatten(x, actx))
+ r_lpot = unflatten(x, actx.from_numpy(mat @ y), actx)
+
+ # }}}
+
+ # {{{ check matvec
+
+ id_eps = 10.0 ** (-np.arange(2, 16))
+ rec_error = np.zeros_like(id_eps)
+
+ from pytools.convergence import EOCRecorder
+ eoc = EOCRecorder()
+
+ from pytential.linalg.skeletonization import skeletonize_by_proxy
+ for i in range(id_eps.size):
+ skeleton = skeletonize_by_proxy(
+ actx, places, tgt_src_index, sym_op, sym_u,
+ domains=[dd], context={},
+ approx_nproxy=case.proxy_approx_count,
+ proxy_radius_factor=case.proxy_radius_factor,
+ id_eps=id_eps[i],
+ )
+ r_hmat = hmatrix_matvec_single_level(mat, y, skeleton[0, 0])
+ r_hmat = unflatten(x, actx.from_numpy(r_hmat), actx)
+
+ from meshmode.dof_array import flat_norm
+ rec_error[i] = actx.to_numpy(
+ flat_norm(r_hmat - r_lpot) / flat_norm(r_lpot)
+ )
+ logger.info("id_eps %.2e error: %.12e", id_eps[i], rec_error[i])
+ # assert rec_error[i] < 0.1
+
+ eoc.add_data_point(id_eps[i], rec_error[i])
+
+ logger.info("\n%s", eoc.pretty_print(
+ abscissa_format="%.8e",
+ error_format="%.8e",
+ eoc_format="%.2f"))
+
+ # }}}
+
+ if not visualize:
+ return
+
+ import matplotlib.pyplot as pt
+ fig = pt.figure(figsize=(10, 10), dpi=300)
+ ax = fig.gca()
+
+ ax.loglog(id_eps, id_eps, "k--")
+ ax.loglog(id_eps, rec_error)
+
+ ax.grid(True)
+ ax.set_xlabel(r"$\epsilon_{id}$")
+ ax.set_ylabel("$Error$")
+ ax.set_title(case.name)
+
+ basename = "linalg_hmatrix_single_matvec"
+ fig.savefig(f"{basename}_{case.name}_{case.op_type}_convergence")
+
+ if case.ambient_dim == 2:
+ fig.clf()
+ ax = fig.gca()
+
+ from arraycontext import flatten
+ r_hmap = actx.to_numpy(flatten(r_hmat, actx))
+ r_lpot = actx.to_numpy(flatten(r_lpot, actx))
+
+ ax.semilogy(r_hmap - r_lpot)
+ ax.set_ylim([1.0e-16, 1.0])
+ fig.savefig(f"{basename}_{case.name}_{case.op_type}_error")
+
+ pt.close(fig)
+
+# }}}
+
+
+# {{{ test_hmatrix_forward_matvec
+
+@pytest.mark.parametrize("case", [
+ HMATRIX_TEST_CASES[0],
+ HMATRIX_TEST_CASES[1],
+ pytest.param(HMATRIX_TEST_CASES[2], marks=pytest.mark.slowtest),
+ ])
+@pytest.mark.parametrize("discr_stage", [
+ sym.QBX_SOURCE_STAGE1,
+ # sym.QBX_SOURCE_STAGE2
+ ])
+def test_hmatrix_forward_matvec(
+ actx_factory, case, discr_stage, p2p=False, visualize=False):
+ actx = actx_factory()
+
+ if visualize:
+ logging.basicConfig(level=logging.INFO)
+
+ if case.ambient_dim == 2:
+ kwargs = {"proxy_approx_count": 64, "proxy_radius_factor": 1.25}
+ else:
+ kwargs = {"proxy_approx_count": 256, "proxy_radius_factor": 1.25}
+
+ case = replace(case, skel_discr_stage=discr_stage, **kwargs)
+ 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[-1], case.target_order)
+ places = GeometryCollection(qbx, auto_where=dd)
+
+ density_discr = places.get_discretization(dd.geometry, dd.discr_stage)
+ max_particles_in_box = case.max_particles_in_box_for_discr(density_discr)
+
+ tgt_src_index, _ = case.get_tgt_src_cluster_index(
+ actx, places, dd, max_particles_in_box=max_particles_in_box)
+
+ 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_by_proxy
+ sym_u, sym_op = case.get_operator(places.ambient_dim)
+
+ x = actx.thaw(density_discr.nodes()[0])
+
+ if p2p:
+ # NOTE: this also needs changed in `build_hmatrix_by_proxy`
+ # to actually evaluate the p2p interactions instead of qbx
+ from pytential.linalg.direct_solver_symbolic import prepare_expr
+ from pytential.symbolic.matrix import P2PMatrixBuilder
+ mat = P2PMatrixBuilder(
+ actx,
+ dep_expr=sym_u,
+ other_dep_exprs=[],
+ dep_discr=density_discr,
+ places=places,
+ context={},
+ )(prepare_expr(places, sym_op, (dd, dd)))
+
+ from arraycontext import flatten, unflatten
+ y = actx.to_numpy(flatten(x, actx))
+ r_lpot = unflatten(x, actx.from_numpy(mat @ y), actx)
+ else:
+ r_lpot = bind(places, sym_op, auto_where=dd)(actx, u=x)
+
+ from pytential.linalg.hmatrix import hmatrix_error_from_param
+ id_eps = 10.0 ** (-np.arange(2, 16))
+ rec_error = np.zeros_like(id_eps)
+ model_error = np.zeros_like(id_eps)
+
+ from pytools.convergence import EOCRecorder
+ eoc = EOCRecorder()
+
+ for i in range(id_eps.size):
+ wrangler = build_hmatrix_by_proxy(actx, places, sym_op, sym_u,
+ domains=[dd],
+ context=case.knl_concrete_kwargs,
+ id_eps=id_eps[i],
+ _tree_kind=case.tree_kind,
+ _max_particles_in_box=max_particles_in_box,
+ _approx_nproxy=case.proxy_approx_count,
+ _proxy_radius_factor=case.proxy_radius_factor,
+ )
+ hmat = wrangler.get_forward()
+
+ # {{{ skeletonization error
+
+ from meshmode.dof_array import flat_norm
+ r_hmap = hmat @ x
+ rec_error[i] = actx.to_numpy(
+ flat_norm(r_hmap - r_lpot) / flat_norm(r_lpot)
+ )
+
+ # }}}
+
+ # {{{ model error
+
+ skeleton = hmat.skeletons[0]
+ icluster = np.argmax(np.diff(skeleton.skel_tgt_src_index.targets.starts))
+
+ proxy_radius = actx.to_numpy(
+ skeleton._src_eval_result.pxy.radii[icluster]
+ )
+ cluster_radius = actx.to_numpy(
+ skeleton._src_eval_result.pxy._cluster_radii[icluster]
+ )
+
+ model_error[i] = hmatrix_error_from_param(
+ places.ambient_dim,
+ id_eps=id_eps[i],
+ min_proxy_radius=proxy_radius,
+ max_cluster_radius=cluster_radius,
+ id_rank=skeleton.skel_tgt_src_index.targets.cluster_size(icluster),
+ nproxies=skeleton._src_eval_result.pxy.pxyindex.cluster_size(icluster),
+ ntargets=skeleton.tgt_src_index.targets.cluster_size(icluster),
+ nsources=skeleton.tgt_src_index.targets.cluster_size(icluster),
+ c=1.0e-8
+ )
+
+ # }}}
+
+ logger.info("id_eps %.2e error: %.12e (%.12e)",
+ id_eps[i], rec_error[i], model_error[i])
+ eoc.add_data_point(id_eps[i], rec_error[i])
+
+ logger.info("\n%s", eoc.pretty_print(
+ abscissa_format="%.8e",
+ error_format="%.8e",
+ eoc_format="%.2f"))
+
+ if not visualize:
+ assert eoc.order_estimate() > 0.6
+
+ # }}}
+
+ if not visualize:
+ return
+
+ import matplotlib.pyplot as pt
+ fig = pt.figure(figsize=(10, 10), dpi=300)
+ ax = fig.gca()
+
+ ax.loglog(id_eps, id_eps, "k--")
+ ax.loglog(id_eps, rec_error)
+ ax.loglog(id_eps, model_error)
+
+ ax.grid(True)
+ ax.set_xlabel(r"$\epsilon_{id}$")
+ ax.set_ylabel("$Error$")
+ ax.set_title(case.name)
+
+ lpot_name = "p2p" if p2p else "qbx"
+ basename = f"linalg_hmatrix_{lpot_name}_matvec"
+ fig.savefig(f"{basename}_{case.name}_{case.op_type}_convergence")
+
+ if case.ambient_dim == 2:
+ fig.clf()
+ 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.semilogy(r_hmap - r_lpot)
+ ax.set_ylim([1.0e-16, 1.0])
+ fig.savefig(f"{basename}_{case.name}_{case.op_type}_error")
+
+ pt.close(fig)
+
+# }}}
+
+
+# {{{ test_hmatrix_backward_matvec
+
+@pytest.mark.parametrize("case", [
+ HMATRIX_TEST_CASES[0],
+ HMATRIX_TEST_CASES[1],
+ pytest.param(HMATRIX_TEST_CASES[2], marks=pytest.mark.slowtest),
+ ])
+@pytest.mark.parametrize("discr_stage", [
+ sym.QBX_SOURCE_STAGE1,
+ # sym.QBX_SOURCE_STAGE2
+ ])
+def test_hmatrix_backward_matvec(actx_factory, case, discr_stage, visualize=False):
+ actx = actx_factory()
+
+ if visualize:
+ logging.basicConfig(level=logging.INFO)
+
+ if case.ambient_dim == 2:
+ kwargs = {"proxy_approx_count": 64, "proxy_radius_factor": 1.25}
+ else:
+ kwargs = {"proxy_approx_count": 64, "proxy_radius_factor": 1.25}
+
+ case = replace(case, skel_discr_stage=discr_stage, **kwargs)
+ 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[-1], case.target_order)
+ places = GeometryCollection(qbx, auto_where=dd)
+
+ density_discr = places.get_discretization(dd.geometry, dd.discr_stage)
+ max_particles_in_box = case.max_particles_in_box_for_discr(density_discr)
+
+ tgt_src_index, _ = case.get_tgt_src_cluster_index(
+ actx, places, dd, max_particles_in_box=max_particles_in_box)
+
+ logger.info("dd %s", dd)
+ logger.info("nclusters %3d ndofs %7d",
+ tgt_src_index.nclusters, density_discr.ndofs)
+
+ # }}}
+
+ # {{{
+
+ sym_u, sym_op = case.get_operator(places.ambient_dim)
+
+ if visualize:
+ from pytential.linalg.direct_solver_symbolic import prepare_expr
+ from pytential.symbolic.matrix import MatrixBuilder
+ mat = MatrixBuilder(
+ actx,
+ dep_expr=sym_u,
+ other_dep_exprs=[],
+ dep_discr=density_discr,
+ places=places,
+ context={},
+ )(prepare_expr(places, sym_op, (dd, dd)))
+
+ import pytential.linalg.utils as hla
+ eigs_ref = hla.eigs(mat, k=5)
+ kappa_ref = np.linalg.cond(mat, p=2)
+
+ # }}}
+
+ # {{{ construct hmatrix
+
+ from pytential.linalg.hmatrix import build_hmatrix_by_proxy
+ sym_u, sym_op = case.get_operator(places.ambient_dim)
+
+ x_ref = actx.thaw(density_discr.nodes()[0])
+ b_ref = bind(places, sym_op, auto_where=dd)(actx, u=x_ref)
+
+ id_eps = 10.0 ** (-np.arange(2, 16))
+ rec_error = np.zeros_like(id_eps)
+
+ if visualize:
+ rec_eigs = np.zeros((id_eps.size, eigs_ref.size), dtype=np.complex128)
+ rec_kappa = np.zeros(id_eps.size)
+
+ from pytools.convergence import EOCRecorder
+ eoc = EOCRecorder()
+
+ for i in range(id_eps.size):
+ wrangler = build_hmatrix_by_proxy(actx, places, sym_op, sym_u,
+ domains=[dd],
+ context=case.knl_concrete_kwargs,
+ id_eps=id_eps[i],
+ _tree_kind=case.tree_kind,
+ _max_particles_in_box=max_particles_in_box,
+ _approx_nproxy=case.proxy_approx_count,
+ _proxy_radius_factor=case.proxy_radius_factor,
+ )
+
+ hmat_inv = wrangler.get_backward()
+ x_hmat = hmat_inv @ b_ref
+
+ if visualize:
+ hmat = wrangler.get_forward()
+ rec_eigs[i, :] = hla.eigs(hmat, k=5, tol=1.0e-6)
+ rec_kappa[i] = hla.cond(hmat, p=2, tol=1.0e-6)
+
+ logger.info("eigs: %s %s", eigs_ref, rec_eigs[i])
+ logger.info("kappa %.12e %.12e", kappa_ref, rec_kappa[i])
+
+ from meshmode.dof_array import flat_norm
+ rec_error[i] = actx.to_numpy(
+ flat_norm(x_hmat - x_ref) / flat_norm(x_ref)
+ )
+ logger.info("id_eps %.2e error: %.12e", id_eps[i], rec_error[i])
+ eoc.add_data_point(id_eps[i], rec_error[i])
+
+ logger.info("\n%s", eoc.pretty_print(
+ abscissa_format="%.8e",
+ error_format="%.8e",
+ eoc_format="%.2f"))
+
+ if not visualize:
+ assert eoc.order_estimate() > 0.6
+
+ # }}}
+
+ if not visualize:
+ return
+
+ import matplotlib.pyplot as pt
+ fig = pt.figure(figsize=(10, 10), dpi=300)
+
+ # {{{ convergence
+
+ ax = fig.gca()
+ ax.loglog(id_eps, id_eps, "k--")
+ ax.loglog(id_eps, rec_error)
+
+ ax.grid(True)
+ ax.set_xlabel(r"$\epsilon_{id}$")
+ ax.set_ylabel("$Error$")
+ ax.set_title(case.name)
+
+ fig.savefig(f"linalg_hmatrix_inverse_{case.name}_{case.op_type}_convergence")
+ fig.clf()
+
+ # }}}
+
+ # {{{ eigs
+
+ ax = fig.gca()
+ ax.plot(np.real(eigs_ref), np.imag(eigs_ref), "ko")
+ for i in range(id_eps.size):
+ ax.plot(np.real(rec_eigs[i]), np.imag(rec_eigs[i]), "v")
+
+ ax.grid(True)
+ ax.set_xlabel(r"$\Re \lambda$")
+ ax.set_ylabel(r"$\Im \lambda$")
+
+ fig.savefig(f"linalg_hmatrix_inverse_{case.name}_{case.op_type}_eigs")
+ fig.clf()
+
+ # }}}
+
+ if case.ambient_dim == 2:
+ ax = fig.gca()
+
+ from arraycontext import flatten
+ x_hmat = actx.to_numpy(flatten(x_hmat, actx))
+ x_ref = actx.to_numpy(flatten(x_ref, actx))
+
+ ax.semilogy(x_hmat - x_ref)
+ ax.set_ylim([1.0e-16, 1.0])
+ fig.savefig(f"linalg_hmatrix_inverse_{case.name}_{case.op_type}_error")
+ fig.clf()
+
+ pt.close(fig)
+
+# }}}
+
+
+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_skeletonization.py b/test/test_linalg_skeletonization.py
index b7be1060f..04da07dea 100644
--- a/test/test_linalg_skeletonization.py
+++ b/test/test_linalg_skeletonization.py
@@ -138,6 +138,105 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False):
# }}}
+# {{{ test_skeletonize_diagonal
+
+@pytest.mark.parametrize("case", [
+ SKELETONIZE_TEST_CASES[0],
+ SKELETONIZE_TEST_CASES[1],
+ SKELETONIZE_TEST_CASES[2],
+ ])
+def test_skeletonize_diagonal(actx_factory, case, visualize=False):
+ import scipy.linalg.interpolative as sli # pylint:disable=no-name-in-module
+ sli.seed(42)
+
+ actx = actx_factory()
+
+ if visualize:
+ logging.basicConfig(level=logging.INFO)
+
+ # {{{ setup
+
+ dd = sym.DOFDescriptor(case.name, discr_stage=case.skel_discr_stage)
+ resolution = case.resolutions[-1]
+
+ qbx = case.get_layer_potential(actx, resolution, case.target_order)
+ places = GeometryCollection(qbx, auto_where=dd)
+
+ tgt_src_index, ctree = case.get_tgt_src_cluster_index(actx, places, dd)
+
+ sym_u, sym_op = case.get_operator(places.ambient_dim)
+
+ # }}}
+
+ # {{{ check
+
+ from pytential.linalg.skeletonization import make_skeletonization_wrangler
+ wrangler = make_skeletonization_wrangler(
+ places, sym_op, sym_u,
+ auto_where=dd, context=case.knl_concrete_kwargs)
+
+ from pytential.linalg.skeletonization import rec_skeletonize_by_proxy
+ skeletons = rec_skeletonize_by_proxy(
+ actx, places, ctree, tgt_src_index, sym_op, sym_u,
+ auto_where=dd,
+ context=case.knl_concrete_kwargs,
+ approx_nproxy=case.proxy_approx_count,
+ proxy_radius_factor=case.proxy_radius_factor,
+ id_eps=case.id_eps,
+ _wrangler=wrangler,
+ )
+
+ from pytential.linalg.hmatrix import _update_skeleton_diagonal
+ for i in range(1, skeletons.size):
+ skeletons[i] = _update_skeleton_diagonal(
+ skeletons[i], skeletons[i - 1], ctree.levels[i - 1],
+ )
+
+ from pytential.linalg.cluster import cluster
+ parent = None
+ for k, clevel in enumerate(ctree.levels):
+ from pytential.linalg.utils import make_flat_cluster_diag
+
+ tgt_src_index = skeletons[k].tgt_src_index
+ D1 = wrangler.evaluate_self(actx, places, tgt_src_index, 0, 0)
+ D1 = make_flat_cluster_diag(D1, tgt_src_index)
+
+ if k == 0:
+ D = D0 = D1
+ else:
+ skel_tgt_src_index = skeletons[k - 1].skel_tgt_src_index
+ assert skel_tgt_src_index.shape == tgt_src_index.shape
+
+ D0 = wrangler.evaluate_self(actx, places, skel_tgt_src_index, 0, 0)
+ D0 = cluster(make_flat_cluster_diag(D0, skel_tgt_src_index), parent)
+
+ D = D1 - D0
+
+ parent = clevel
+
+ assert D1.shape == (skeletons[k].nclusters,)
+ assert D1.shape == D0.shape, (D1.shape, D0.shape)
+ assert D1.shape == D.shape, (D1.shape, D.shape)
+
+ for i in range(skeletons[k].nclusters):
+ assert D1[i].shape == skeletons[k].tgt_src_index.cluster_shape(i, i)
+ assert D1[i].shape == D0[i].shape, (D1[i].shape, D0[i].shape)
+
+ error = la.norm(D[i] - skeletons[k].D[i]) / (la.norm(D[i]) + 1.0e-12)
+ logger.info("level %04d / %04d cluster %3d (%4d, %4d) error %.12e",
+ k, ctree.nlevels,
+ ctree.tree_cluster_parent_ids[clevel.box_ids][i],
+ *skeletons[k].tgt_src_index.cluster_shape(i, i), error)
+
+ assert error < 1.0e-15
+
+ logger.info("")
+
+ # }}}
+
+# }}}
+
+
# {{{ test_skeletonize_by_proxy