Direct Solver: Skeletonization

pytential/linalg/__init__.py

@@ -23,18 +23,27 @@
 from pytential.linalg.utils import (
         IndexList, TargetAndSourceClusterList,
         make_index_list, make_index_cluster_cartesian_product,
+        interp_decomp,
         )
 from pytential.linalg.proxy import (
-        ProxyClusterGeometryData,
+        ProxyClusterGeometryData, ProxyPointTarget, ProxyPointSource,
         ProxyGeneratorBase, ProxyGenerator, QBXProxyGenerator,
         partition_by_nodes, gather_cluster_neighbor_points,
         )
+from pytential.linalg.skeletonization import (
+        SkeletonizationWrangler, make_skeletonization_wrangler,
+        SkeletonizationResult, skeletonize_by_proxy,
+        )
 
 __all__ = (
     "IndexList", "TargetAndSourceClusterList",
     "make_index_list", "make_index_cluster_cartesian_product",
+    "interp_decomp",
 
-    "ProxyClusterGeometryData",
+    "ProxyClusterGeometryData", "ProxyPointTarget", "ProxyPointSource",
     "ProxyGeneratorBase", "ProxyGenerator", "QBXProxyGenerator",
     "partition_by_nodes", "gather_cluster_neighbor_points",
+
+    "SkeletonizationWrangler", "make_skeletonization_wrangler",
+    "SkeletonizationResult", "skeletonize_by_proxy",
 )

pytential/linalg/direct_solver_symbolic.py

@@ -20,6 +20,8 @@
 THE SOFTWARE.
 """
 
+from pytools.obj_array import make_obj_array
+
 from pytential.symbolic.mappers import (
         IdentityMapper, OperatorCollector, LocationTagger)
 
@@ -30,6 +32,39 @@
 """
 
 
+# {{{ utils
+
+class PROXY_SKELETONIZATION_SOURCE:             # noqa: N801
+    pass
+
+
+class PROXY_SKELETONIZATION_TARGET:             # noqa: N801
+    pass
+
+
+def prepare_expr(places, exprs, auto_where=None):
+    from pytential.symbolic.execution import _prepare_expr
+    return make_obj_array([
+        _prepare_expr(places, expr, auto_where=auto_where)
+        for expr in exprs])
+
+
+def prepare_proxy_expr(places, exprs, auto_where=None):
+    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)
+        # ensure all IntGs have their source and targets set
+        expr = DOFDescriptorReplacer(auto_where[0], auto_where[1])(expr)
+
+        return expr
+
+    return make_obj_array([_prepare_expr(expr) for expr in exprs])
+
+# }}}
+
+
 # {{{ KernelTransformationRemover
 
 class KernelTransformationRemover(IdentityMapper):

pytential/linalg/proxy.py

@@ -33,6 +33,9 @@
 from pytential import GeometryCollection, bind, sym
 from pytential.symbolic.dof_desc import DOFDescriptorLike
 from pytential.linalg.utils import IndexList
+from pytential.source import PointPotentialSource
+from pytential.target import PointsTarget
+from pytential.qbx import QBXLayerPotentialSource
 
 import loopy as lp
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
@@ -44,7 +47,10 @@
 
 .. currentmodule:: pytential.linalg
 
+.. autoclass:: ProxyPointSource
+.. autoclass:: ProxyPointTarget
 .. autoclass:: ProxyClusterGeometryData
+
 .. autoclass:: ProxyGeneratorBase
 .. autoclass:: ProxyGenerator
 .. autoclass:: QBXProxyGenerator
@@ -122,28 +128,49 @@ def partition_by_nodes(
 # }}}
 
 
-# {{{ proxy point generator
-
-def _generate_unit_sphere(ambient_dim: int, approx_npoints: int) -> np.ndarray:
-    """Generate uniform points on a unit sphere.
+# {{{ proxy points
 
-    :arg ambient_dim: dimension of the ambient space.
-    :arg approx_npoints: approximate number of points to generate. If the
-        ambient space is 3D, this will not generate the exact number of points.
-    :return: array of shape ``(ambient_dim, npoints)``, where ``npoints``
-        will not generally be the same as ``approx_npoints``.
+class ProxyPointSource(PointPotentialSource):
+    """
+    .. automethod:: get_expansion_for_qbx_direct_eval
     """
 
-    if ambient_dim == 2:
-        t = np.linspace(0.0, 2.0 * np.pi, approx_npoints)
-        points = np.vstack([np.cos(t), np.sin(t)])
-    elif ambient_dim == 3:
-        from pytools import sphere_sample_equidistant
-        points = sphere_sample_equidistant(approx_npoints, r=1)
-    else:
-        raise ValueError("ambient_dim > 3 not supported.")
+    def __init__(self,
+            lpot_source: QBXLayerPotentialSource,
+            proxies: np.ndarray) -> None:
+        """
+        :arg lpot_source: the layer potential for which the proxy are constructed.
+        :arg proxies: an array of shape ``(ambient_dim, nproxies)`` containing
+            the proxy points.
+        """
+        assert proxies.shape[0] == lpot_source.ambient_dim
+
+        super().__init__(proxies)
+        self.lpot_source = lpot_source
+
+    def get_expansion_for_qbx_direct_eval(self, base_kernel, target_kernels):
+        """Wrapper around
+        ``pytential.qbx.QBXLayerPotentialSource.get_expansion_for_qbx_direct_eval``
+        to allow this class to be used by the matrix builders.
+        """
+        return self.lpot_source.get_expansion_for_qbx_direct_eval(
+                base_kernel, target_kernels)
 
-    return points
+
+class ProxyPointTarget(PointsTarget):
+    def __init__(self,
+            lpot_source: QBXLayerPotentialSource,
+            proxies: np.ndarray) -> None:
+        """
+        :arg lpot_source: the layer potential for which the proxy are constructed.
+            This argument is kept for symmetry with :class:`ProxyPointSource`.
+        :arg proxies: an array of shape ``(ambient_dim, nproxies)`` containing
+            the proxy points.
+        """
+        assert proxies.shape[0] == lpot_source.ambient_dim
+
+        super().__init__(proxies)
+        self.lpot_source = lpot_source
 
 
 @dataclass(frozen=True)
@@ -176,6 +203,8 @@ class ProxyClusterGeometryData:
 
     .. automethod:: __init__
     .. automethod:: to_numpy
+    .. automethod:: as_sources
+    .. automethod:: as_targets
     """
 
     places: GeometryCollection
@@ -188,17 +217,65 @@ class ProxyClusterGeometryData:
     centers: np.ndarray
     radii: np.ndarray
 
+    _cluster_radii: Optional[np.ndarray] = None
+
     @property
     def nclusters(self) -> int:
         return self.srcindex.nclusters
 
+    @property
+    def discr(self):
+        return self.places.get_discretization(
+                self.dofdesc.geometry,
+                self.dofdesc.discr_stage)
+
     def to_numpy(self, actx: PyOpenCLArrayContext) -> "ProxyClusterGeometryData":
         from arraycontext import to_numpy
+        if self._cluster_radii is not None:
+            cluster_radii = to_numpy(self._cluster_radii, actx)
+        else:
+            cluster_radii = None
+
         from dataclasses import replace
         return replace(self,
-                points=to_numpy(self.points, actx),
-                centers=to_numpy(self.centers, actx),
-                radii=to_numpy(self.radii, actx))
+                points=np.stack(to_numpy(self.points, actx)),
+                centers=np.stack(to_numpy(self.centers, actx)),
+                radii=to_numpy(self.radii, actx),
+                _cluster_radii=cluster_radii)
+
+    def as_sources(self) -> ProxyPointSource:
+        lpot_source = self.places.get_geometry(self.dofdesc.geometry)
+        return ProxyPointSource(lpot_source, self.points)
+
+    def as_targets(self) -> ProxyPointTarget:
+        lpot_source = self.places.get_geometry(self.dofdesc.geometry)
+        return ProxyPointTarget(lpot_source, self.points)
+
+# }}}
+
+
+# {{{ proxy point generator
+
+def _generate_unit_sphere(ambient_dim: int, approx_npoints: int) -> np.ndarray:
+    """Generate uniform points on a unit sphere.
+
+    :arg ambient_dim: dimension of the ambient space.
+    :arg approx_npoints: approximate number of points to generate. If the
+        ambient space is 3D, this will not generate the exact number of points.
+    :return: array of shape ``(ambient_dim, npoints)``, where ``npoints``
+        will not generally be the same as ``approx_npoints``.
+    """
+
+    if ambient_dim == 2:
+        t = np.linspace(0.0, 2.0 * np.pi, approx_npoints, endpoint=False)
+        points = np.vstack([np.cos(t), np.sin(t)])
+    elif ambient_dim == 3:
+        from pytools import sphere_sample_equidistant
+        points = sphere_sample_equidistant(approx_npoints, r=1)
+    else:
+        raise ValueError("ambient_dim > 3 not supported.")
+
+    return points
 
 
 def make_compute_cluster_centers_knl(
@@ -346,6 +423,8 @@ def __call__(self,
         discr = self.places.get_discretization(
                 source_dd.geometry, source_dd.discr_stage)
 
+        include_cluster_radii = kwargs.pop("include_cluster_radii", False)
+
         # {{{ get proxy centers and radii
 
         sources = flatten(discr.nodes(), actx, leaf_class=DOFArray)
@@ -365,6 +444,18 @@ def __call__(self,
                 proxy_centers=centers_dev,
                 **kwargs)
 
+        if include_cluster_radii:
+            knl = make_compute_cluster_radii_knl(actx, self.ambient_dim)
+            _, (cluster_radii,) = knl(actx.queue,
+                    sources=sources,
+                    srcindices=dof_index.indices,
+                    srcstarts=dof_index.starts,
+                    radius_factor=1.0,
+                    proxy_centers=centers_dev)
+            cluster_radii = actx.freeze(cluster_radii)
+        else:
+            cluster_radii = None
+
         # }}}
 
         # {{{ build proxy points for each cluster
@@ -398,6 +489,7 @@ def __call__(self,
                 points=actx.freeze(from_numpy(proxies, actx)),
                 centers=actx.freeze(centers_dev),
                 radii=actx.freeze(radii_dev),
+                _cluster_radii=cluster_radii
                 )