Import and update BlockIndexRanges from sumpy

doc/linalg.rst

@@ -4,6 +4,12 @@ Linear Algebra Routines
 Hierarchical Direct Solver
 --------------------------
 
+.. warning::
+
+    All the classes and routines in this module are experimental and the
+    API can change at any point.
+
 .. automodule:: pytential.linalg.proxy
+.. automodule:: pytential.linalg.utils
 
 .. vim: sw=4:tw=75

pytential/linalg/__init__.py

@@ -19,3 +19,22 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
+
+from pytential.linalg.proxy import (
+        BlockProxyPoints, ProxyGeneratorBase,
+        ProxyGenerator, QBXProxyGenerator,
+        partition_by_nodes, gather_block_neighbor_points,
+        )
+from pytential.linalg.utils import (
+        BlockIndexRanges, MatrixBlockIndexRanges,
+        make_block_index_from_array, make_index_blockwise_product,
+        )
+
+__all__ = (
+    "BlockProxyPoints", "ProxyGeneratorBase",
+    "ProxyGenerator", "QBXProxyGenerator",
+    "partition_by_nodes", "gather_block_neighbor_points",
+
+    "BlockIndexRanges", "MatrixBlockIndexRanges",
+    "make_block_index_from_array", "make_index_blockwise_product",
+)

pytential/linalg/proxy.py

@@ -30,7 +30,7 @@
 from meshmode.discretization import Discretization
 
 from pytools import memoize_in
-from sumpy.tools import BlockIndexRanges
+from pytential.linalg.utils import BlockIndexRanges
 
 import loopy as lp
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
@@ -40,6 +40,8 @@
 Proxy Point Generation
 ~~~~~~~~~~~~~~~~~~~~~~
 
+.. currentmodule:: pytential.linalg
+
 .. autoclass:: BlockProxyPoints
 .. autoclass:: ProxyGeneratorBase
 .. autoclass:: ProxyGenerator
@@ -52,24 +54,6 @@
 
 # {{{ point index partitioning
 
-def make_block_index(
-        actx: PyOpenCLArrayContext,
-        indices: np.ndarray,
-        ranges: Optional[np.ndarray] = None) -> BlockIndexRanges:
-    """Wrap a ``(indices, ranges)`` tuple into a ``BlockIndexRanges``.
-
-    :param ranges: if *None*, then *indices* is expected to be an object
-        array of indices, so that the ranges can be reconstructed.
-    """
-    if ranges is None:
-        ranges = np.cumsum([0] + [r.size for r in indices])
-        indices = np.hstack(indices)
-
-    return BlockIndexRanges(actx.context,
-            actx.freeze(actx.from_numpy(indices)),
-            actx.freeze(actx.from_numpy(ranges)))
-
-
 def partition_by_nodes(
         actx: PyOpenCLArrayContext, discr: Discretization, *,
         tree_kind: Optional[str] = "adaptive-level-restricted",
@@ -118,7 +102,8 @@ def partition_by_nodes(
         ranges = np.linspace(0, discr.ndofs, nblocks + 1, dtype=np.int64)
         assert ranges[-1] == discr.ndofs
 
-    return make_block_index(actx, indices, ranges=ranges)
+    from pytential.linalg import make_block_index_from_array
+    return make_block_index_from_array(indices, ranges=ranges)
 
 # }}}
 
@@ -176,12 +161,12 @@ class BlockProxyPoints:
     """
     .. attribute:: srcindices
 
-        A :class:`~sumpy.tools.BlockIndexRanges` describing which block of
+        A :class:`~pytential.linalg.BlockIndexRanges` describing which block of
         points each proxy ball was created from.
 
     .. attribute:: indices
 
-        A :class:`~sumpy.tools.BlockIndexRanges` describing which proxies
+        A :class:`~pytential.linalg.BlockIndexRanges` describing which proxies
         belong to which block.
 
     .. attribute:: points
@@ -220,8 +205,6 @@ def to_numpy(self, actx: PyOpenCLArrayContext) -> "BlockProxyPoints":
         from arraycontext import to_numpy
         from dataclasses import replace
         return replace(self,
-                srcindices=self.srcindices.get(actx.queue),
-                indices=self.indices.get(actx.queue),
                 points=to_numpy(self.points, actx),
                 centers=to_numpy(self.centers, actx),
                 radii=to_numpy(self.radii, actx))
@@ -390,9 +373,10 @@ def __call__(self,
         pxyranges = np.arange(0, nproxy + 1, self.nproxy)
 
         from arraycontext import freeze, from_numpy
+        from pytential.linalg import make_block_index_from_array
         return BlockProxyPoints(
                 srcindices=indices,
-                indices=make_block_index(actx, pxyindices, pxyranges),
+                indices=make_block_index_from_array(pxyindices, pxyranges),
                 points=freeze(from_numpy(proxies, actx), actx),
                 centers=freeze(centers_dev, actx),
                 radii=freeze(radii_dev, actx),
@@ -610,7 +594,7 @@ def prg():
     from arraycontext import to_numpy
     pxycenters = to_numpy(pxy.centers, actx)
     pxyradii = to_numpy(pxy.radii, actx)
-    indices = pxy.srcindices.get(actx.queue)
+    indices = pxy.srcindices
 
     nbrindices = np.empty(indices.nblocks, dtype=object)
     for iproxy in range(indices.nblocks):
@@ -639,7 +623,8 @@ def prg():
 
     # }}}
 
-    return make_block_index(actx, indices=nbrindices)
+    from pytential.linalg import make_block_index_from_array
+    return make_block_index_from_array(indices=nbrindices)
 
 # }}}