implement hierachical matrix solver

Author: alexfikl

.gitignore

@@ -21,6 +21,8 @@ examples/*.pdf
 
 *.vtu
 *.vts
+*.png
+*.gif
 
 .cache
 

doc/conf.py

@@ -35,6 +35,7 @@
 }
 
 nitpick_ignore_regex = [
+    ["py:class", r".*_ProxyNeighborEvaluationResult"],
     # Sphinx started complaining about these in 8.2.1(-ish)
     # -AK, 2025-02-24
     ["py:class", r"TypeAliasForwardRef"],

examples/scaling-study-hmatrix.py

@@ -0,0 +1,199 @@
+__copyright__ = "Copyright (C) 2022 Alexandru Fikl"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+import logging
+from dataclasses import dataclass
+
+import numpy as np
+
+from meshmode.array_context import PyOpenCLArrayContext
+from pytools.convergence import EOCRecorder
+
+from pytential import GeometryCollection, sym
+
+
+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)
+
+    eoc_build = EOCRecorder()
+    eoc_matvec = EOCRecorder()
+
+    import meshmode.discretization.poly_element as mpoly
+    import meshmode.mesh.generation as mgen
+
+    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)

pytential/linalg/cluster.py

@@ -208,7 +208,7 @@ def cluster(obj: object, clevel: ClusterLevel) -> Any:
 
 
 @cluster.register(IndexList)
-def _cluster_index_list(obj: IndexList, clevel: ClusterLevel) -> IndexList:
+def cluster_index_list(obj: IndexList, clevel: ClusterLevel) -> IndexList:
     assert obj.nclusters == clevel.nclusters
 
     if clevel.nclusters == 1:
@@ -224,7 +224,7 @@ def _cluster_index_list(obj: IndexList, clevel: ClusterLevel) -> IndexList:
 
 
 @cluster.register(TargetAndSourceClusterList)
-def _cluster_target_and_source_cluster_list(
+def cluster_target_and_source_cluster_list(
         obj: TargetAndSourceClusterList, clevel: ClusterLevel,
         ) -> TargetAndSourceClusterList:
     assert obj.nclusters == clevel.nclusters
@@ -238,7 +238,7 @@ def _cluster_target_and_source_cluster_list(
 
 
 @cluster.register(np.ndarray)
-def _cluster_ndarray(obj: NDArray[Any], clevel: ClusterLevel) -> NDArray[Any]:
+def cluster_ndarray(obj: NDArray[Any], clevel: ClusterLevel) -> NDArray[Any]:
     assert obj.shape == (clevel.nclusters,)
     if clevel.nclusters == 1:
         return obj
@@ -420,26 +420,14 @@ def partition_by_nodes(
 
 # {{{ visualize clusters
 
-def visualize_clusters(actx: PyOpenCLArrayContext,
-                       generator: ProxyGenerator,
-                       srcindex: IndexList,
-                       tree: ClusterTree,
-                       filename: str | pathlib.Path, *,
-                       dofdesc: sym.DOFDescriptorLike = None,
-                       overwrite: bool = False) -> None:
-    filename = pathlib.Path(filename)
-
-    places = generator.places
-    if dofdesc is None:
-        dofdesc = places.auto_source
-    dofdesc = sym.as_dofdesc(dofdesc)
-
-    discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage)
-    assert isinstance(discr, Discretization)
-
-    if discr.ambient_dim != 2:
-        raise NotImplementedError(f"Unsupported dimension: {discr.ambient_dim}")
-
+def _visualize_clusters_2d(actx: PyOpenCLArrayContext,
+                           discr: Discretization,
+                           generator: ProxyGenerator,
+                           srcindex: IndexList,
+                           tree: ClusterTree,
+                           filename: pathlib.Path, *,
+                           dofdesc: sym.DOFDescriptor,
+                           overwrite: bool = False) -> None:
     import matplotlib.pyplot as pt
 
     from arraycontext import flatten
@@ -448,7 +436,7 @@ def visualize_clusters(actx: PyOpenCLArrayContext,
 
     x, y = actx.to_numpy(flatten(discr.nodes(), actx, leaf_class=DOFArray))
     for clevel in tree.levels:
-        outfile = filename.with_stem(f"{filename.stem}-{clevel.level:03d}")
+        outfile = filename.with_stem(f"{filename.stem}-lvl{clevel.level:03d}")
         if not overwrite and outfile.exists():
             raise FileExistsError(f"Output file '{outfile}' already exists")
 
@@ -493,6 +481,102 @@ def visualize_clusters(actx: PyOpenCLArrayContext,
 
         srcindex = cluster(srcindex, clevel)
 
+
+def _visualize_clusters_3d(actx: PyOpenCLArrayContext,
+                           discr: Discretization,
+                           generator: ProxyGenerator,
+                           srcindex: IndexList,
+                           tree: ClusterTree,
+                           filename: pathlib.Path, *,
+                           dofdesc: sym.DOFDescriptor,
+                           overwrite: bool = False) -> None:
+    from arraycontext import unflatten
+    from meshmode.discretization.visualization import make_visualizer
+
+    for clevel in tree.levels:
+        outfile = filename.with_stem(f"{filename.stem}-lvl{clevel.level:03d}")
+        outfile = outfile.with_suffix(".vtu")
+        if not overwrite and outfile.exists():
+            raise FileExistsError(f"Output file '{outfile}' already exists")
+
+        # construct proxy balls
+        pxy = generator(actx, dofdesc, srcindex).to_numpy(actx)
+        pxycenters = pxy.centers
+        pxyradii = pxy.radii
+        nclusters = srcindex.nclusters
+
+        # construct meshes for each proxy ball
+        from meshmode.mesh.generation import generate_sphere
+        from meshmode.mesh.processing import affine_map, merge_disjoint_meshes
+
+        ref_mesh = generate_sphere(1, 4, uniform_refinement_rounds=1)
+        pxymeshes = [
+            affine_map(ref_mesh, A=pxyradii[i], b=pxycenters[:, i].squeeze())
+            for i in range(nclusters)
+        ]
+
+        # merge meshes into a single discretization
+        from meshmode.discretization.poly_element import (
+            InterpolatoryEdgeClusteredGroupFactory,
+        )
+        pxymesh = merge_disjoint_meshes([discr.mesh, *pxymeshes])
+        pxydiscr = Discretization(actx, pxymesh,
+                                  InterpolatoryEdgeClusteredGroupFactory(4))
+
+        # add a marker field for all clusters
+        marker = np.full((pxydiscr.ndofs,), np.nan, dtype=np.float64)
+        template_ary = actx.thaw(pxydiscr.nodes()[0])
+
+        for i in range(srcindex.nclusters):
+            isrc = srcindex.cluster_indices(i)
+            marker[isrc] = 10.0 * (i + 1.0)
+        marker_dev = unflatten(template_ary, actx.from_numpy(marker), actx)
+
+        # add a marker field for all proxies
+        pxymarker = np.full((pxydiscr.ndofs,), np.nan, dtype=np.float64)
+        pxymarker[discr.ndofs:] = 1.0
+        pxymarker_dev = unflatten(template_ary, actx.from_numpy(pxymarker), actx)
+
+        # write it all out
+        vis = make_visualizer(actx, pxydiscr)
+        vis.write_vtk_file(outfile, [
+            ("marker", marker_dev),
+            ("proxies", pxymarker_dev),
+            ], overwrite=overwrite)
+
+        srcindex = cluster(srcindex, clevel)
+
+
+def visualize_clusters(actx: PyOpenCLArrayContext,
+                       generator: ProxyGenerator,
+                       srcindex: IndexList,
+                       tree: ClusterTree,
+                       filename: str | pathlib.Path, *,
+                       dofdesc: sym.DOFDescriptorLike = None,
+                       overwrite: bool = False) -> None:
+    filename = pathlib.Path(filename)
+
+    places = generator.places
+    if dofdesc is None:
+        dofdesc = places.auto_source
+    dofdesc = sym.as_dofdesc(dofdesc)
+
+    discr = places.get_discretization(dofdesc.geometry, dofdesc.discr_stage)
+    assert isinstance(discr, Discretization)
+
+    if discr.ambient_dim == 2:
+        _visualize_clusters_2d(
+            actx, discr, generator, srcindex, tree, filename,
+            dofdesc=dofdesc,
+            overwrite=overwrite)
+    elif discr.ambient_dim == 3:
+        _visualize_clusters_3d(
+            actx, discr, generator, srcindex, tree, filename,
+            dofdesc=dofdesc,
+            overwrite=overwrite)
+    else:
+        raise NotImplementedError(f"Unsupported dimension: {discr.ambient_dim}")
+
 # }}}