From 43fdf44112510c473bbe58cb5f97daea975bac41 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 10:42:09 +0300
Subject: [PATCH 01/10] docs: fix missing TreeKind reference

---
 doc/conf.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/doc/conf.py b/doc/conf.py
index 71043e18a..770ac2a47 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -40,7 +40,6 @@
     ["py:class", r"TypeAliasForwardRef"],
     ["py:class", r"arraycontext.typing._UserDefinedArrayContainer"],
     ["py:class", r"arraycontext.typing._UserDefinedArithArrayContainer"],
-    ["py:class", r"TreeKind"],
 ]
 
 
@@ -70,7 +69,7 @@
     "Discretization": "class:meshmode.discretization.Discretization",
     "DOFArray": "class:meshmode.dof_array.DOFArray",
     # boxtree
-    # "TreeKind": "obj:boxtree.tree_build.TreeKind",
+    "TreeKind": "obj:boxtree.tree_build.TreeKind",
     # sumpy
     "ExpansionBase": "class:sumpy.expansion.ExpansionBase",
     "ExpansionFactoryBase": "class:sumpy.expansion.ExpansionFactoryBase",

From ed1537fba0525db7921d6a61220b12fe970184cf Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 11:05:35 +0300
Subject: [PATCH 02/10] docs: improve docs for linalg.proxy

---
 pytential/linalg/direct_solver_symbolic.py |  2 +
 pytential/linalg/proxy.py                  | 93 ++++++++++++----------
 2 files changed, 55 insertions(+), 40 deletions(-)

diff --git a/pytential/linalg/direct_solver_symbolic.py b/pytential/linalg/direct_solver_symbolic.py
index e0cbb24f7..2e478af8c 100644
--- a/pytential/linalg/direct_solver_symbolic.py
+++ b/pytential/linalg/direct_solver_symbolic.py
@@ -261,3 +261,5 @@ def __init__(self, source: DOFDescriptorLike, target: DOFDescriptorLike) -> None
         self.operand_rec = _LocationReplacer(source, default_source=source)
 
 # }}}
+
+# vim: foldmethod=marker
diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py
index 8d98191ad..db77fcde8 100644
--- a/pytential/linalg/proxy.py
+++ b/pytential/linalg/proxy.py
@@ -32,7 +32,6 @@
 
 import loopy as lp
 from arraycontext import Array, ArrayContainer, PyOpenCLArrayContext, flatten
-from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 from meshmode.discretization import Discretization
 from meshmode.dof_array import DOFArray
 from pytools import memoize_in
@@ -208,30 +207,15 @@ def __init__(self,
 @dataclass(frozen=True)
 class ProxyClusterGeometryData:
     """
-    .. attribute:: srcindex
+    .. autoattribute:: places
+    .. autoattribute:: dofdesc
 
-        A :class:`~pytential.linalg.IndexList` describing which cluster
-        of points each proxy ball was created from.
-
-    .. attribute:: pxyindex
-
-        A :class:`~pytential.linalg.IndexList` describing which proxies
-        belong to which cluster.
-
-    .. attribute:: points
-
-        A concatenated array of all the proxy points. Can be sliced into
-        using :attr:`pxyindex` (shape ``(dim, nproxies)``).
-
-    .. attribute:: centers
-
-        An array of all the proxy ball centers (shape ``(dim, nclusters)``).
-
-    .. attribute:: radii
-
-        An array of all the proxy ball radii (shape ``(nclusters,)``).
-
-    .. attribute:: nclusters
+    .. autoattribute:: srcindex
+    .. autoattribute:: pxyindex
+    .. autoattribute:: points
+    .. autoattribute:: centers
+    .. autoattribute:: radii
+    .. autoattribute:: nclusters
 
     .. automethod:: __init__
     .. automethod:: to_numpy
@@ -241,22 +225,36 @@ class ProxyClusterGeometryData:
 
     places: GeometryCollection
     dofdesc: sym.DOFDescriptor
+    """A descriptor for the geometry used to compute the proxy points."""
 
     srcindex: IndexList
+    """A :class:`~pytential.linalg.IndexList` describing which cluster
+    of points each proxy ball was created from.
+    """
     pxyindex: IndexList
+    """A :class:`~pytential.linalg.IndexList` describing which proxies
+    belong to which cluster.
+    """
 
     points: Array
+    """A concatenated array of all the proxy points. Can be sliced into
+    using :attr:`pxyindex` (shape ``(dim, nproxies)``).
+    """
     centers: Array
+    """An array of all the proxy ball centers (shape ``(dim, nclusters)``)."""
     radii: Array
+    """An array of all the proxy ball radii (shape ``(nclusters,)``)."""
 
     _cluster_radii: Array | None = None
 
     @property
     def nclusters(self) -> int:
+        """Number of proxy clusters."""
         return self.srcindex.nclusters
 
     @property
     def discr(self) -> Discretization:
+        """The discretization that corresponds to :attr:`dofdesc` in :attr:`places`."""
         discr = self.places.get_discretization(
             self.dofdesc.geometry, self.dofdesc.discr_stage)
         assert isinstance(discr, Discretization)
@@ -361,7 +359,7 @@ def prg() -> lp.ExecutorBase:
             name="compute_cluster_centers_knl",
             assumptions="ndim>=1 and nclusters>=1",
             fixed_parameters={"ndim": ndim},
-            lang_version=MOST_RECENT_LANGUAGE_VERSION,
+            lang_version=lp.MOST_RECENT_LANGUAGE_VERSION,
             )
 
         knl = lp.tag_inames(knl, "idim*:unr")
@@ -374,9 +372,12 @@ def prg() -> lp.ExecutorBase:
 
 class ProxyGeneratorBase:
     r"""
-    .. attribute:: nproxy
+    .. autoattribute:: places
+    .. autoattribute:: radius_factor
+    .. autoattribute:: norm_type
+    .. autoattribute:: ref_points
+    .. autoattribute:: nproxy
 
-        Number of proxy points in a single proxy ball.
 
     .. automethod:: __init__
     .. automethod:: __call__
@@ -384,8 +385,18 @@ class ProxyGeneratorBase:
 
     places: GeometryCollection
     radius_factor: float
+    """Factor multiplying the cluster radius. This is currently fixed for all
+    clusters.
+    """
     norm_type: str
+    """The type of the norm used to compute the bounding box of the cluster. The
+    "radius" is also computed in terms of this norm. Can be one of ``"l2"`` or
+    ``"linf"``.
+    """
     ref_points: NDArray[np.floating]
+    """Reference proxy points on the unit sphere. The reference points are
+    translated and scaled for each cluster.
+    """
 
     def __init__(
             self,
@@ -436,10 +447,12 @@ def __init__(
 
     @property
     def ambient_dim(self) -> int:
+        """Ambient dimension of the proxy geometry."""
         return self.places.ambient_dim
 
     @property
     def nproxy(self) -> int:
+        """Number of proxy points in a single proxy ball."""
         return self.ref_points.shape[1]
 
     def get_centers_kernel_ex(self, actx: PyOpenCLArrayContext) -> lp.ExecutorBase:
@@ -464,10 +477,9 @@ def __call__(self,
         """
         if source_dd is None:
             source_dd = self.places.auto_source
-        source_dd = sym.as_dofdesc(source_dd)
+        dd = sym.as_dofdesc(source_dd)
 
-        discr = self.places.get_discretization(
-                source_dd.geometry, source_dd.discr_stage)
+        discr = self.places.get_discretization(dd.geometry, dd.discr_stage)
         assert isinstance(discr, Discretization)
 
         # {{{ get proxy centers and radii
@@ -526,7 +538,7 @@ def __call__(self,
         from pytential.linalg import make_index_list
         return ProxyClusterGeometryData(
                 places=self.places,
-                dofdesc=source_dd,
+                dofdesc=dd,
                 srcindex=dof_index,
                 pxyindex=make_index_list(pxyindices, pxystarts),
                 points=actx.freeze(actx.from_numpy(proxies)),
@@ -566,7 +578,7 @@ def prg():
             name="compute_cluster_radii_knl",
             assumptions="ndim>=1 and nclusters>=1",
             fixed_parameters={"ndim": ndim},
-            lang_version=MOST_RECENT_LANGUAGE_VERSION,
+            lang_version=lp.MOST_RECENT_LANGUAGE_VERSION,
             )
 
         knl = lp.tag_inames(knl, "idim*:unr")
@@ -631,7 +643,7 @@ def prg() -> lp.ExecutorBase:
             name="compute_cluster_qbx_radii_knl",
             assumptions="ndim>=1 and nclusters>=1",
             fixed_parameters={"ndim": ndim},
-            lang_version=MOST_RECENT_LANGUAGE_VERSION,
+            lang_version=lp.MOST_RECENT_LANGUAGE_VERSION,
             )
 
         knl = lp.tag_inames(knl, "idim*:unr")
@@ -662,16 +674,16 @@ def __call__(self,
             **kwargs: ArrayContainer) -> ProxyClusterGeometryData:
         if source_dd is None:
             source_dd = self.places.auto_source
-        source_dd = sym.as_dofdesc(source_dd)
+        dd = sym.as_dofdesc(source_dd)
 
         radii = cast("ArrayContainer", bind(self.places, sym.expansion_radii(
-            self.ambient_dim, dofdesc=source_dd))(actx))
+            self.ambient_dim, dofdesc=dd))(actx))
         center_int = cast("ArrayContainer", bind(self.places, sym.expansion_centers(
-            self.ambient_dim, -1, dofdesc=source_dd))(actx))
+            self.ambient_dim, -1, dofdesc=dd))(actx))
         center_ext = cast("ArrayContainer", bind(self.places, sym.expansion_centers(
-            self.ambient_dim, +1, dofdesc=source_dd))(actx))
+            self.ambient_dim, +1, dofdesc=dd))(actx))
 
-        return super().__call__(actx, source_dd, dof_index,
+        return super().__call__(actx, dd, dof_index,
                 expansion_radii=flatten(radii, actx),
                 center_int=flatten(center_int, actx, leaf_class=DOFArray),
                 center_ext=flatten(center_ext, actx, leaf_class=DOFArray),
@@ -684,7 +696,8 @@ def __call__(self,
 # {{{ gather_cluster_neighbor_points
 
 def gather_cluster_neighbor_points(
-        actx: PyOpenCLArrayContext, pxy: ProxyClusterGeometryData, *,
+        actx: PyOpenCLArrayContext,
+        pxy: ProxyClusterGeometryData, *,
         max_particles_in_box: int | None = None) -> IndexList:
     """Generate a set of neighboring points for each cluster of points in
     *pxy*. Neighboring points of a cluster :math:`i` are defined
@@ -721,7 +734,7 @@ def prg() -> lp.ExecutorBase:
             name="picker_knl",
             assumptions="ndim>=1 and npoints>=1",
             fixed_parameters={"ndim": discr.ambient_dim},
-            lang_version=MOST_RECENT_LANGUAGE_VERSION,
+            lang_version=lp.MOST_RECENT_LANGUAGE_VERSION,
             )
 
         knl = lp.tag_inames(knl, "idim*:unr")

From 7ddb53c8cfc0c181a9a4fe574a70c82500611008 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 11:15:39 +0300
Subject: [PATCH 03/10] docs: improve docs for linalg.utils

---
 pytential/linalg/utils.py | 70 ++++++++++++++++++++++++---------------
 1 file changed, 44 insertions(+), 26 deletions(-)

diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py
index d7bc12898..33325cf84 100644
--- a/pytential/linalg/utils.py
+++ b/pytential/linalg/utils.py
@@ -24,15 +24,18 @@
 """
 
 from dataclasses import dataclass
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypeVar
 
 import numpy as np
 import numpy.linalg as la
 
-from pytools import memoize_in, memoize_method, obj_array
+from pytools import memoize_in, obj_array
 
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     from numpy.typing import NDArray
 
     from arraycontext import Array, PyOpenCLArrayContext
@@ -64,31 +67,37 @@
 class IndexList:
     """Convenience class for working with clusters (subsets) of an array.
 
-    .. attribute:: nclusters
-    .. attribute:: indices
-
-        An :class:`~numpy.ndarray` of not necessarily continuous or increasing
-        integers representing the indices of a global array. The individual
-        cluster slices are delimited using :attr:`starts`.
-
-    .. attribute:: starts
-
-        An :class:`~numpy.ndarray` of size ``(nclusters + 1,)`` consisting of
-        nondecreasing integers used to index into :attr:`indices`. A cluster
-        :math:`i` can be retrieved using ``indices[starts[i]:starts[i + 1]]``.
-
+    .. autoattribute:: nclusters
+    .. autoattribute:: nindices
+    .. autoattribute:: indices
+    .. autoattribute:: starts
     .. automethod:: cluster_size
     .. automethod:: cluster_indices
     .. automethod:: cluster_take
     """
 
     indices: NDArray[np.integer]
+    """An :class:`~numpy.ndarray` of not necessarily continuous or increasing
+    integers representing the indices of a global array. The individual
+    cluster slices are delimited using :attr:`starts`.
+    """
+
     starts: NDArray[np.integer]
+    """An :class:`~numpy.ndarray` of size ``(nclusters + 1,)`` consisting of
+    nondecreasing integers used to index into :attr:`indices`. A cluster
+    :math:`i` can be retrieved using ``indices[starts[i]:starts[i + 1]]``.
+    """
 
     @property
     def nclusters(self) -> int:
+        """Number of clusters in the index list."""
         return self.starts.size - 1
 
+    @property
+    def nindices(self) -> int:
+        """Total number of indices in the index list."""
+        return self.indices.size
+
     def cluster_size(self, i: int | np.integer) -> int:
         if not (0 <= i < self.nclusters):
             raise IndexError(
@@ -125,14 +134,10 @@ def cluster_take(self,
 class TargetAndSourceClusterList:
     """Convenience class for working with clusters (subsets) of a matrix.
 
-    .. attribute:: nclusters
-    .. attribute:: targets
-
-        An :class:`IndexList` encapsulating target cluster indices.
-
-    .. attribute:: sources
-
-        An :class:`IndexList` encapsulating source cluster indices.
+    .. autoattribute:: nclusters
+    .. autoattribute:: shape
+    .. autoattribute:: targets
+    .. autoattribute:: sources
 
     .. automethod:: cluster_shape
     .. automethod:: cluster_indices
@@ -141,7 +146,9 @@ class TargetAndSourceClusterList:
     """
 
     targets: IndexList
+    """An :class:`IndexList` encapsulating target cluster indices."""
     sources: IndexList
+    """An :class:`IndexList` encapsulating source cluster indices."""
 
     def __post_init__(self) -> None:
         if self.targets.nclusters != self.sources.nclusters:
@@ -152,10 +159,18 @@ def __post_init__(self) -> None:
 
     @property
     def nclusters(self) -> int:
+        """Number of clusters in the index list."""
         return self.targets.nclusters
 
     @property
-    @memoize_method
+    def shape(self) -> tuple[int, int]:
+        """Shape of the Cartesian product of the :attr:`targets` and :attr:`sources`."""
+        return (self.targets.nindices, self.sources.nindices)
+
+    def __iter__(self) -> Iterator[IndexList]:
+        return iter((self.targets, self.sources))
+
+    @cached_property
     def _flat_cluster_starts(self) -> NDArray[np.integer]:
         return np.cumsum([0] + [
             self.targets.cluster_size(i) * self.sources.cluster_size(i)
@@ -356,8 +371,9 @@ def interp_decomp(
     :return: a tuple ``(k, idx, interp)`` containing the numerical rank *k*,
         the column indices *idx* and the resulting interpolation matrix *interp*.
     """
-    if rank is not None and eps is not None:
-        raise ValueError("providing both 'rank' and 'eps' is not supported")
+
+    if (rank is not None and eps is not None) or (rank is None and eps is None):
+        raise ValueError("either 'rank' or 'eps' must be provided (not both)")
 
     from scipy import __version__
 
@@ -461,7 +477,7 @@ def mnorm(x: NDArray[Any], y: NDArray[Any]) -> np.floating[Any]:
 
 
 def skeletonization_error(
-        mat: np.ndarray, skeleton: SkeletonizationResult, *,
+        mat: NDArray[Any], skeleton: SkeletonizationResult, *,
         ord: float | None = None,
         relative: bool = False) -> np.floating[Any]:
     r"""Computes the skeletonization error for the entire matrix *mat*.
@@ -520,3 +536,5 @@ def skeletonization_error(
     return result
 
 # }}}
+
+# vim: foldmethod=marker

From 6f6bad43cd7dd32a3dea1a1f2da674a3b53be59f Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 11:17:02 +0300
Subject: [PATCH 04/10] feat: make make_flat_cluster_diag return a 1d array

---
 pytential/linalg/utils.py | 11 +++++------
 test/test_matrix.py       |  2 +-
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py
index 33325cf84..8e082be3e 100644
--- a/pytential/linalg/utils.py
+++ b/pytential/linalg/utils.py
@@ -337,20 +337,19 @@ def _product() -> tuple[Array, Array]:
 def make_flat_cluster_diag(
         mat: NDArray[InexactT],
         mindex: TargetAndSourceClusterList,
-    ) -> obj_array.ObjectArray2D[NDArray[InexactT]]:
+    ) -> obj_array.ObjectArray1D[NDArray[InexactT]]:
     """
     :param mat: a one-dimensional :class:`~numpy.ndarray` that has a one-to-one
         correspondence to the index sets constructed by
         :func:`make_index_cluster_cartesian_product` for *mindex*.
 
-    :returns: a block diagonal object :class:`~numpy.ndarray`, where each
-        diagonal element :math:`(i, i)` is the reshaped slice of *mat* that
-        corresponds to the cluster :math:`i`.
+    :returns: an object :class:`~numpy.ndarray`, where each element represents
+        the block of a block-diagonal matrix.
     """
-    cluster_mat = np.full((mindex.nclusters, mindex.nclusters), 0, dtype=object)
+    cluster_mat = np.empty(mindex.nclusters, dtype=object)
     for i in range(mindex.nclusters):
         shape = mindex.cluster_shape(i, i)
-        cluster_mat[i, i] = mindex.flat_cluster_take(mat, i).reshape(*shape)
+        cluster_mat[i] = mindex.flat_cluster_take(mat, i).reshape(*shape)
 
     return cluster_mat
 
diff --git a/test/test_matrix.py b/test/test_matrix.py
index 221205814..97f27d38c 100644
--- a/test/test_matrix.py
+++ b/test/test_matrix.py
@@ -63,7 +63,7 @@ def max_cluster_error(mat, clusters, mindex, p=None):
 
         error = max(
                 error,
-                la.norm(mat_i - clusters[i, i], ord=p) / norm_mat_i
+                la.norm(mat_i - clusters[i], ord=p) / norm_mat_i
                 )
 
     return error

From 17cdcbfa9c2b4f1b875ef141c048cc9020a1fbda Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 11:29:07 +0300
Subject: [PATCH 05/10] test: small tweaks to linalg skeletonization tests

---
 test/test_linalg_proxy.py           |  7 ++---
 test/test_linalg_skeletonization.py | 41 ++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/test/test_linalg_proxy.py b/test/test_linalg_proxy.py
index b61871625..17c3e7b3b 100644
--- a/test/test_linalg_proxy.py
+++ b/test/test_linalg_proxy.py
@@ -38,13 +38,10 @@
 
 from pytential import GeometryCollection, bind, sym
 from pytential.linalg import ProxyGenerator, QBXProxyGenerator
-
-
-logger = logging.getLogger(__name__)
-
 from pytential.utils import pytest_teardown_function as teardown_function  # noqa: F401
 
 
+logger = logging.getLogger(__name__)
 pytest_generate_tests = pytest_generate_tests_for_array_contexts([
     PytestPyOpenCLArrayContextFactory,
     ])
@@ -240,7 +237,7 @@ def test_partition_points(actx_factory, tree_kind, case, visualize=False):
     ProxyGenerator, QBXProxyGenerator,
     ])
 @pytest.mark.parametrize("index_sparsity_factor", [1.0, 0.6])
-@pytest.mark.parametrize("proxy_radius_factor", [1, 1.1])
+@pytest.mark.parametrize("proxy_radius_factor", [1.0, 1.1])
 def test_proxy_generator(actx_factory, case,
         proxy_generator_cls, index_sparsity_factor, proxy_radius_factor,
         visualize=False):
diff --git a/test/test_linalg_skeletonization.py b/test/test_linalg_skeletonization.py
index 1f91379be..f1015e7ad 100644
--- a/test/test_linalg_skeletonization.py
+++ b/test/test_linalg_skeletonization.py
@@ -114,6 +114,7 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False):
     """
 
     actx = actx_factory()
+    rng = np.random.default_rng(42)
 
     if visualize:
         logging.basicConfig(level=logging.INFO)
@@ -154,7 +155,6 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False):
 
     from pytential.linalg.skeletonization import _skeletonize_block_by_proxy_with_mats
 
-    rng = np.random.default_rng(42)
     _skeletonize_block_by_proxy_with_mats(
         actx, 0, 0, places, proxy_generator, wrangler, tgt_src_index,
         id_eps=1.0e-8,
@@ -170,6 +170,7 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False):
 def run_skeletonize_by_proxy(actx, case, resolution,
                              places=None, mat=None,
                              ctol=None, rtol=None,
+                             tgt_src_index=None,
                              rng=None,
                              suffix="", visualize=False):
     from pytools import ProcessTimer
@@ -180,9 +181,12 @@ def run_skeletonize_by_proxy(actx, case, resolution,
     if places is None:
         qbx = case.get_layer_potential(actx, resolution, case.target_order)
         places = GeometryCollection(qbx, auto_where=dd)
+    else:
+        qbx = places.get_geometry(dd.geometry)
 
     density_discr = places.get_discretization(dd.geometry, dd.discr_stage)
-    tgt_src_index = case.get_tgt_src_cluster_index(actx, places, dd)
+    if tgt_src_index is None:
+        tgt_src_index = case.get_tgt_src_cluster_index(actx, places, dd)
 
     logger.info("nclusters %3d ndofs %7d",
             tgt_src_index.nclusters, density_discr.ndofs)
@@ -207,8 +211,8 @@ def run_skeletonize_by_proxy(actx, case, resolution,
 
     sym_u, sym_op = case.get_operator(places.ambient_dim)
     wrangler = make_skeletonization_wrangler(places, sym_op, sym_u,
-            domains=None,
             context=case.knl_concrete_kwargs,
+            auto_where=dd,
             _weighted_proxy=case.weighted_proxy,
             _proxy_source_cluster_builder=case.proxy_source_cluster_builder,
             _proxy_target_cluster_builder=case.proxy_target_cluster_builder,
@@ -252,10 +256,13 @@ def run_skeletonize_by_proxy(actx, case, resolution,
 
     logger.info("[time] skeletonization by proxy: %s", p)
 
+    def intersect1d(x, y):
+        return np.where((x.reshape(1, -1) - y.reshape(-1, 1)) == 0)[1]
+
     L, R = skeleton.L, skeleton.R
     for i in range(tgt_src_index.nclusters):
         # targets (rows)
-        bi = np.searchsorted(
+        bi = intersect1d(
             tgt_src_index.targets.cluster_indices(i),
             skeleton.skel_tgt_src_index.targets.cluster_indices(i),
             )
@@ -265,7 +272,7 @@ def run_skeletonize_by_proxy(actx, case, resolution,
         tgt_error = la.norm(A - L[i] @ S, ord=ord) / la.norm(A, ord=ord)
 
         # sources (columns)
-        bj = np.searchsorted(
+        bj = intersect1d(
             tgt_src_index.sources.cluster_indices(i),
             skeleton.skel_tgt_src_index.sources.cluster_indices(i),
             )
@@ -279,8 +286,8 @@ def run_skeletonize_by_proxy(actx, case, resolution,
                 src_error, tgt_error, R[i].shape[0], R[i].shape[1])
 
         if ctol is not None:
-            assert src_error < ctol * case.id_eps
-            assert tgt_error < ctol * case.id_eps
+            assert src_error < ctol
+            assert tgt_error < ctol
 
     # }}}
 
@@ -313,9 +320,9 @@ def run_skeletonize_by_proxy(actx, case, resolution,
             rtol if rtol is not None else 0.0)
 
     if rtol:
-        assert err_l < rtol * case.id_eps
-        assert err_r < rtol * case.id_eps
-        assert err_f < rtol * case.id_eps
+        assert err_l < rtol
+        assert err_r < rtol
+        assert err_f < rtol
 
     # }}}
 
@@ -385,9 +392,9 @@ def test_skeletonize_by_proxy(actx_factory, case, visualize=False):
 
     run_skeletonize_by_proxy(
         actx, case, case.resolutions[0],
-        ctol=10,
+        ctol=10 * case.id_eps,
         # FIXME: why is the 3D error so large?
-        rtol=10**case.ambient_dim,
+        rtol=10**case.ambient_dim * case.id_eps,
         rng=rng,
         visualize=visualize)
 
@@ -455,9 +462,13 @@ def test_skeletonize_by_proxy_convergence(
     places = mat = None
     for i in range(id_eps.size):
         case = replace(case, id_eps=id_eps[i], weighted_proxy=weighted)
-        rec_error[i], (places, mat) = run_skeletonize_by_proxy(
-            actx, case, r, places=places, mat=mat,
-            suffix=f"{suffix}_{i:04d}", rng=rng, visualize=False)
+
+        # NOTE: don't skeletonize anymore if we reached zero error, but we still
+        # want to loop to do `eoc.add_data_point()`
+        if not was_zero:
+            rec_error[i], (places, mat) = run_skeletonize_by_proxy(
+                actx, case, r, places=places, mat=mat,
+                suffix=f"{suffix}_{i:04d}", rng=rng, visualize=False)
 
         was_zero = rec_error[i] == 0.0
         eoc.add_data_point(id_eps[i], rec_error[i])

From 936732ae07855afff371eef131aecd2005663d22 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 15:24:38 +0300
Subject: [PATCH 06/10] feat: do not export everything in linalg.__init__

---
 doc/conf.py                         |  7 ++++-
 doc/linalg.rst                      | 21 ++++++++++++---
 pytential/linalg/__init__.py        | 23 ----------------
 pytential/linalg/proxy.py           |  8 +++---
 pytential/linalg/skeletonization.py | 41 ++++++++++++++---------------
 pytential/linalg/utils.py           |  3 ---
 pytential/symbolic/matrix.py        | 11 ++++----
 test/extra_matrix_data.py           |  6 ++---
 test/test_linalg_proxy.py           |  4 +--
 test/test_linalg_skeletonization.py |  8 +++---
 10 files changed, 61 insertions(+), 71 deletions(-)

diff --git a/doc/conf.py b/doc/conf.py
index 770ac2a47..03a4b5f75 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -38,8 +38,9 @@
     # Sphinx started complaining about these in 8.2.1(-ish)
     # -AK, 2025-02-24
     ["py:class", r"TypeAliasForwardRef"],
-    ["py:class", r"arraycontext.typing._UserDefinedArrayContainer"],
+    ["py:class", r"_ProxyNeighborEvaluationResult"],
     ["py:class", r"arraycontext.typing._UserDefinedArithArrayContainer"],
+    ["py:class", r"arraycontext.typing._UserDefinedArrayContainer"],
 ]
 
 
@@ -48,10 +49,13 @@
     "NDArray": "obj:numpy.typing.NDArray",
     "np.integer": "obj:numpy.integer",
     "np.floating": "obj:numpy.floating",
+    "np.inexact": "obj:numpy.inexact",
+    "np.random.Generator": "class:numpy.random.Generator",
     # pytools
     "ObjectArrayND": "obj:pytools.obj_array.ObjectArrayND",
     "T": "obj:pytools.T",
     "obj_array.ObjectArray1D": "obj:pytools.obj_array.ObjectArray1D",
+    "obj_array.ObjectArray2D": "obj:pytools.obj_array.ObjectArray2D",
     # pyopencl
     "WaitList": "obj:pyopencl.WaitList",
     # pymbolic
@@ -87,6 +91,7 @@
         "data:pytential.symbolic.dof_desc.DOFDescriptorLike",
     "sym.DOFDescriptor": "class:pytential.symbolic.dof_desc.DOFDescriptor",
     "sym.IntG": "class:pytential.symbolic.primitives.IntG",
+    "sym.var": "obj:pytential.symbolic.primitives.var",
 }
 
 
diff --git a/doc/linalg.rst b/doc/linalg.rst
index 97497434f..de26425a0 100644
--- a/doc/linalg.rst
+++ b/doc/linalg.rst
@@ -8,7 +8,7 @@ scheme is used:
   component of the Stokeslet.
 * ``cluster`` refers to a piece of a ``block`` as used by the recursive
   proxy-based skeletonization of the direct solver algorithms. Clusters
-  are represented by a :class:`~pytential.linalg.TargetAndSourceClusterList`.
+  are represented by a :class:`~pytential.linalg.utils.TargetAndSourceClusterList`.
 
 GMRES
 -----
@@ -20,17 +20,30 @@ GMRES
 Hierarchical Direct Solver
 --------------------------
 
+.. note::
+
+    High-level API for direct solvers is in progress.
+
+Low-level Functionality
+-----------------------
+
 .. 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
+.. automodule:: pytential.linalg.skeletonization
+
+Internal Functionality and Utilities
+------------------------------------
 
-Internal Functionality
-----------------------
+.. warning::
 
+    All the classes and routines in this module are experimental and the
+    API can change at any point.
+
+.. automodule:: pytential.linalg.utils
 .. automodule:: pytential.linalg.direct_solver_symbolic
 
 .. vim: sw=4:tw=75:fdm=marker
diff --git a/pytential/linalg/__init__.py b/pytential/linalg/__init__.py
index f3b8c45e0..1869399ce 100644
--- a/pytential/linalg/__init__.py
+++ b/pytential/linalg/__init__.py
@@ -23,22 +23,6 @@
 THE SOFTWARE.
 """
 
-from pytential.linalg.proxy import (
-        ProxyClusterGeometryData,
-        ProxyGenerator,
-        ProxyGeneratorBase,
-        ProxyPointSource,
-        ProxyPointTarget,
-        QBXProxyGenerator,
-        gather_cluster_neighbor_points,
-        partition_by_nodes,
-)
-from pytential.linalg.skeletonization import (
-        SkeletonizationResult,
-        SkeletonizationWrangler,
-        make_skeletonization_wrangler,
-        skeletonize_by_proxy,
-)
 from pytential.linalg.utils import (
         IndexList,
         TargetAndSourceClusterList,
@@ -52,11 +36,4 @@
     "IndexList", "TargetAndSourceClusterList",
     "make_index_list", "make_index_cluster_cartesian_product",
     "interp_decomp",
-
-    "ProxyClusterGeometryData", "ProxyPointTarget", "ProxyPointSource",
-    "ProxyGeneratorBase", "ProxyGenerator", "QBXProxyGenerator",
-    "partition_by_nodes", "gather_cluster_neighbor_points",
-
-    "SkeletonizationWrangler", "make_skeletonization_wrangler",
-    "SkeletonizationResult", "skeletonize_by_proxy",
 )
diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py
index db77fcde8..633fadc16 100644
--- a/pytential/linalg/proxy.py
+++ b/pytential/linalg/proxy.py
@@ -59,8 +59,6 @@
 Proxy Point Generation
 ~~~~~~~~~~~~~~~~~~~~~~
 
-.. currentmodule:: pytential.linalg
-
 .. autoclass:: ProxyPointSource
 .. autoclass:: ProxyPointTarget
 .. autoclass:: ProxyClusterGeometryData
@@ -147,7 +145,7 @@ def partition_by_nodes(
 
         assert starts[-1] == discr.ndofs
 
-    from pytential.linalg import make_index_list
+    from pytential.linalg.utils import make_index_list
     return make_index_list(indices, starts=starts)
 
 # }}}
@@ -228,11 +226,11 @@ class ProxyClusterGeometryData:
     """A descriptor for the geometry used to compute the proxy points."""
 
     srcindex: IndexList
-    """A :class:`~pytential.linalg.IndexList` describing which cluster
+    """A :class:`~pytential.linalg.utils.IndexList` describing which cluster
     of points each proxy ball was created from.
     """
     pxyindex: IndexList
-    """A :class:`~pytential.linalg.IndexList` describing which proxies
+    """A :class:`~pytential.linalg.utils.IndexList` describing which proxies
     belong to which cluster.
     """
 
diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py
index ce7942096..155db4b7d 100644
--- a/pytential/linalg/skeletonization.py
+++ b/pytential/linalg/skeletonization.py
@@ -54,8 +54,6 @@
 
 
 __doc__ = """
-.. currentmodule:: pytential.linalg
-
 Skeletonization
 ---------------
 
@@ -126,8 +124,8 @@ def _approximate_geometry_waa_magnitude(
         cluster_index: IndexList,
         domain: sym.DOFDescriptor) -> Array:
     """
-    :arg cluster_index: a :class:`~pytential.linalg.IndexList` representing the
-        clusters on which to operate. In practice, this can be the cluster
+    :arg cluster_index: a :class:`~pytential.linalg.utils.IndexList` representing
+        the clusters on which to operate. In practice, this can be the cluster
         indices themselves or the indices of neighboring points to each
         cluster (i.e. points inside the proxy ball).
 
@@ -508,7 +506,7 @@ class _ProxyNeighborEvaluationResult:
     """
     .. attribute:: pxy
 
-        A :class:`~pytential.linalg.ProxyClusterGeometryData` containing the
+        A :class:`~pytential.linalg.utils.ProxyClusterGeometryData` containing the
         proxy points from which :attr:`pxymat` is obtained. This data is also
         used to construct :attr:`nbrindex` and evaluate :attr:`nbrmat`.
 
@@ -518,12 +516,12 @@ class _ProxyNeighborEvaluationResult:
         target points. This matrix is flattened to a shape of ``(nsize,)``,
         which is consistent with the sum of the cluster sizes in :attr:`pxyindex`,
         as obtained from
-        :meth:`~pytential.linalg.TargetAndSourceClusterList.cluster_size`.
+        :meth:`~pytential.linalg.utils.TargetAndSourceClusterList.cluster_size`.
 
     .. attribute:: pxyindex
 
-        A :class:`~pytential.linalg.TargetAndSourceClusterList` used to describe
-        the cluster interactions in :attr:`pxymat`.
+        A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` used to
+        describe the cluster interactions in :attr:`pxymat`.
 
     .. attribute:: nbrmat
 
@@ -531,12 +529,12 @@ class _ProxyNeighborEvaluationResult:
         target points. This matrix is flattened to a shape of ``(nsize,)``,
         which is consistent with the sum of the cluster sizes in :attr:`nbrindex`,
         as obtained from
-        :meth:`~pytential.linalg.TargetAndSourceClusterList.cluster_size`.
+        :meth:`~pytential.linalg.utils.TargetAndSourceClusterList.cluster_size`.
 
     .. attribute:: nbrindex
 
-        A :class:`~pytential.linalg.TargetAndSourceClusterList` used to describe
-        the cluster interactions in :attr:`nbrmat`.
+        A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` used to
+        describe the cluster interactions in :attr:`nbrmat`.
 
     .. automethod:: __getitem__
     """
@@ -585,7 +583,7 @@ def _evaluate_proxy_skeletonization_interaction(
     if cluster_index.nclusters == 1:
         raise ValueError("cannot make a proxy skeleton for a single cluster")
 
-    from pytential.linalg import gather_cluster_neighbor_points
+    from pytential.linalg.proxy import gather_cluster_neighbor_points
     pxy = proxy_generator(actx, dofdesc, cluster_index)
     nbrindex = gather_cluster_neighbor_points(
             actx, pxy,
@@ -649,7 +647,7 @@ def _skeletonize_block_by_proxy_with_mats(
     L = np.empty(nclusters, dtype=object)
     R = np.empty(nclusters, dtype=object)
 
-    from pytential.linalg import interp_decomp
+    from pytential.linalg.utils import interp_decomp
 
     for i in range(nclusters):
         k = id_rank
@@ -686,7 +684,8 @@ def _skeletonize_block_by_proxy_with_mats(
         skel_starts[i + 1] = skel_starts[i] + k
         assert tgt_skl_indices[i].shape == src_skl_indices[i].shape
 
-    from pytential.linalg import make_index_list
+    from pytential.linalg.utils import make_index_list
+
     src_skl_index = make_index_list(np.hstack(list(src_skl_indices)), skel_starts)
     tgt_skl_index = make_index_list(np.hstack(list(tgt_skl_indices)), skel_starts)
     skel_tgt_src_index = TargetAndSourceClusterList(tgt_skl_index, src_skl_index)
@@ -736,13 +735,13 @@ class SkeletonizationResult:
 
     .. attribute:: tgt_src_index
 
-        A :class:`~pytential.linalg.TargetAndSourceClusterList` representing the
-        indices in the original matrix :math:`A` that have been skeletonized.
+        A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` representing
+        the indices in the original matrix :math:`A` that have been skeletonized.
 
     .. attribute:: skel_tgt_src_index
 
-        A :class:`~pytential.linalg.TargetAndSourceClusterList` representing a
-        subset of :attr:`tgt_src_index`, i.e. the skeleton of each cluster of
+        A :class:`~pytential.linalg.utils.TargetAndSourceClusterList` representing
+        a subset of :attr:`tgt_src_index`, i.e. the skeleton of each cluster of
         :math:`A`. These indices can be used to reconstruct the :math:`S`
         matrix.
     """
@@ -799,7 +798,7 @@ def skeletonize_by_proxy(
     ) -> obj_array.ObjectArray2D[NDArray[Any]]:
     r"""Evaluate and skeletonize a symbolic expression using proxy-based methods.
 
-    :arg tgt_src_index: a :class:`~pytential.linalg.TargetAndSourceClusterList`
+    :arg tgt_src_index: a :class:`~pytential.linalg.utils.TargetAndSourceClusterList`
         indicating which indices participate in the skeletonization.
 
     :arg exprs: see :func:`make_skeletonization_wrangler`.
@@ -807,8 +806,8 @@ def skeletonize_by_proxy(
     :arg domains: see :func:`make_skeletonization_wrangler`.
     :arg context: see :func:`make_skeletonization_wrangler`.
 
-    :arg approx_nproxy: see :class:`~pytential.linalg.ProxyGenerator`.
-    :arg proxy_radius_factor: see :class:`~pytential.linalg.ProxyGenerator`.
+    :arg approx_nproxy: see :class:`~pytential.linalg.proxy.ProxyGenerator`.
+    :arg proxy_radius_factor: see :class:`~pytential.linalg.proxy.ProxyGenerator`.
 
     :arg id_eps: a floating point value used as a tolerance in skeletonizing
         each block in *tgt_src_index*.
diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py
index 8e082be3e..924a3f096 100644
--- a/pytential/linalg/utils.py
+++ b/pytential/linalg/utils.py
@@ -48,9 +48,6 @@
 ~~~~
 
 .. autoclass:: InexactT
-
-.. currentmodule:: pytential.linalg
-
 .. autoclass:: IndexList
 .. autoclass:: TargetAndSourceClusterList
 
diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py
index 86be1b5d3..9f62c5098 100644
--- a/pytential/symbolic/matrix.py
+++ b/pytential/symbolic/matrix.py
@@ -295,7 +295,7 @@ def map_call(self, expr: p.Call):
 
 class ClusterMatrixBuilderBase(MatrixBuilderBase):
     """Evaluate individual clusters of a matrix operator, as defined by a
-    :class:`~pytential.linalg.TargetAndSourceClusterList`.
+    :class:`~pytential.linalg.utils.TargetAndSourceClusterList`.
 
     Unlike, e.g. :class:`MatrixBuilder`, matrix cluster builders are
     significantly reduced in scope. They are basically just meant
@@ -315,7 +315,8 @@ def __init__(self,
                  tgt_src_index: TargetAndSourceClusterList,
                  context: dict[str, Any]) -> None:
         """
-        :arg tgt_src_index: a :class:`~pytential.linalg.TargetAndSourceClusterList`
+        :arg tgt_src_index: a
+            :class:`~pytential.linalg.utils.TargetAndSourceClusterList`
             class describing which clusters are going to be evaluated.
         """
 
@@ -336,7 +337,7 @@ def _inner_mapper(self):
                 self.tgt_src_index, self.context)
 
     def get_dep_variable(self):
-        from pytential.linalg import make_index_cluster_cartesian_product
+        from pytential.linalg.utils import make_index_cluster_cartesian_product
         actx = self.array_context
         tgtindices, srcindices = (
                 make_index_cluster_cartesian_product(actx, self.tgt_src_index)
@@ -679,7 +680,7 @@ def map_int_g(self, expr):
 
             # {{{ geometry
 
-            from pytential.linalg import make_index_cluster_cartesian_product
+            from pytential.linalg.utils import make_index_cluster_cartesian_product
             tgtindices, srcindices = make_index_cluster_cartesian_product(
                     actx, self.tgt_src_index)
 
@@ -786,7 +787,7 @@ def map_int_g(self, expr):
 
             # {{{ geometry
 
-            from pytential.linalg import make_index_cluster_cartesian_product
+            from pytential.linalg.utils import make_index_cluster_cartesian_product
             tgtindices, srcindices = make_index_cluster_cartesian_product(
                     actx, self.tgt_src_index)
 
diff --git a/test/extra_matrix_data.py b/test/extra_matrix_data.py
index dea7b4a38..46ff9686f 100644
--- a/test/extra_matrix_data.py
+++ b/test/extra_matrix_data.py
@@ -14,7 +14,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from pytential.symbolic.dof_desc import DiscretizationStages
+    from pytential.symbolic.dof_desc import DiscretizationStage
 
 
 # {{{ MatrixTestCase
@@ -42,7 +42,7 @@ class MatrixTestCaseMixin:
 
     # skeletonization
     id_eps: float = 1.0e-8
-    skel_discr_stage: DiscretizationStages = sym.QBX_SOURCE_STAGE2
+    skel_discr_stage: DiscretizationStage = sym.QBX_SOURCE_STAGE2
 
     weighted_proxy: bool | None = None
     proxy_source_cluster_builder: Callable[..., Any] | None = None
@@ -58,7 +58,7 @@ def get_cluster_index(self, actx, places, dofdesc=None):
         if max_particles_in_box is None:
             max_particles_in_box = discr.ndofs // self.approx_cluster_count
 
-        from pytential.linalg import partition_by_nodes
+        from pytential.linalg.proxy import partition_by_nodes
         cindex = partition_by_nodes(actx, places,
                 dofdesc=dofdesc,
                 tree_kind=self.tree_kind,
diff --git a/test/test_linalg_proxy.py b/test/test_linalg_proxy.py
index 17c3e7b3b..9115cdfd4 100644
--- a/test/test_linalg_proxy.py
+++ b/test/test_linalg_proxy.py
@@ -37,7 +37,7 @@
 from meshmode.mesh.generation import NArmedStarfish, ellipse
 
 from pytential import GeometryCollection, bind, sym
-from pytential.linalg import ProxyGenerator, QBXProxyGenerator
+from pytential.linalg.proxy import ProxyGenerator, QBXProxyGenerator
 from pytential.utils import pytest_teardown_function as teardown_function  # noqa: F401
 
 
@@ -337,7 +337,7 @@ def test_neighbor_points(actx_factory, case,
     pxy = generator(actx, dofdesc, srcindex)
 
     # get neighboring points
-    from pytential.linalg import gather_cluster_neighbor_points
+    from pytential.linalg.proxy import gather_cluster_neighbor_points
     nbrindex = gather_cluster_neighbor_points(actx, pxy)
 
     pxy = pxy.to_numpy(actx)
diff --git a/test/test_linalg_skeletonization.py b/test/test_linalg_skeletonization.py
index f1015e7ad..675b69362 100644
--- a/test/test_linalg_skeletonization.py
+++ b/test/test_linalg_skeletonization.py
@@ -136,12 +136,12 @@ def test_skeletonize_symbolic(actx_factory, case, visualize=False):
 
     # {{{ wranglers
 
-    from pytential.linalg import QBXProxyGenerator
-    from pytential.linalg.skeletonization import make_skeletonization_wrangler
+    from pytential.linalg.proxy import QBXProxyGenerator
     proxy_generator = QBXProxyGenerator(places,
             radius_factor=case.proxy_radius_factor,
             approx_nproxy=case.proxy_approx_count)
 
+    from pytential.linalg.skeletonization import make_skeletonization_wrangler
     sym_u, sym_op = case.get_operator(places.ambient_dim)
     wrangler = make_skeletonization_wrangler(places, sym_op, sym_u,
             domains=None,
@@ -203,12 +203,12 @@ def run_skeletonize_by_proxy(actx, case, resolution,
     logger.info("proxy factor %.2f count %7d",
                 case.proxy_radius_factor, proxy_approx_count)
 
-    from pytential.linalg import QBXProxyGenerator
-    from pytential.linalg.skeletonization import make_skeletonization_wrangler
+    from pytential.linalg.proxy import QBXProxyGenerator
     proxy_generator = QBXProxyGenerator(places,
             radius_factor=case.proxy_radius_factor,
             approx_nproxy=proxy_approx_count)
 
+    from pytential.linalg.skeletonization import make_skeletonization_wrangler
     sym_u, sym_op = case.get_operator(places.ambient_dim)
     wrangler = make_skeletonization_wrangler(places, sym_op, sym_u,
             context=case.knl_concrete_kwargs,

From a54a7eeb851ef3c644e4cb7dcb7cc87b3af3c143 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 20:24:15 +0300
Subject: [PATCH 07/10] feat: always compute cluster_radii in linalg.proxy

---
 pytential/linalg/proxy.py | 71 ++++++++++++++++++---------------------
 1 file changed, 32 insertions(+), 39 deletions(-)

diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py
index 633fadc16..b7729c8fb 100644
--- a/pytential/linalg/proxy.py
+++ b/pytential/linalg/proxy.py
@@ -23,6 +23,7 @@
 THE SOFTWARE.
 """
 
+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, cast
 
@@ -55,6 +56,9 @@
     from pytential.symbolic.dof_desc import DOFDescriptorLike
 
 
+logger = logging.getLogger(__name__)
+
+
 __doc__ = """
 Proxy Point Generation
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -204,16 +208,18 @@ def __init__(self,
 
 @dataclass(frozen=True)
 class ProxyClusterGeometryData:
-    """
+    """The proxy information generated by  :class:`ProxyGeneratorBase`.
+
     .. autoattribute:: places
     .. autoattribute:: dofdesc
 
+    .. autoattribute:: nclusters
     .. autoattribute:: srcindex
     .. autoattribute:: pxyindex
     .. autoattribute:: points
     .. autoattribute:: centers
     .. autoattribute:: radii
-    .. autoattribute:: nclusters
+    .. autoattribute:: cluster_radii
 
     .. automethod:: __init__
     .. automethod:: to_numpy
@@ -242,12 +248,15 @@ class ProxyClusterGeometryData:
     """An array of all the proxy ball centers (shape ``(dim, nclusters)``)."""
     radii: Array
     """An array of all the proxy ball radii (shape ``(nclusters,)``)."""
-
-    _cluster_radii: Array | None = None
+    cluster_radii: Array
+    """An array of all the cluster ball radii (shape ``(nclusters,)``). The
+    proxy radii are essentially just the cluster radii multiplied by the
+    ``radius_factor``.
+    """
 
     @property
     def nclusters(self) -> int:
-        """Number of proxy clusters."""
+        """Number of clusters."""
         return self.srcindex.nclusters
 
     @property
@@ -260,17 +269,12 @@ def discr(self) -> Discretization:
         return discr
 
     def to_numpy(self, actx: PyOpenCLArrayContext) -> ProxyClusterGeometryData:
-        if self._cluster_radii is not None:
-            cluster_radii = actx.to_numpy(self._cluster_radii)
-        else:
-            cluster_radii = None
-
         from dataclasses import replace
         return replace(self,
                 points=np.stack(actx.to_numpy(self.points)),
                 centers=np.stack(actx.to_numpy(self.centers)),
                 radii=actx.to_numpy(self.radii),
-                _cluster_radii=cluster_radii)
+                cluster_radii=self.cluster_radii)
 
     def as_sources(self) -> ProxyPointSource:
         lpot_source = self.places.get_geometry(self.dofdesc.geometry)
@@ -376,7 +380,6 @@ class ProxyGeneratorBase:
     .. autoattribute:: ref_points
     .. autoattribute:: nproxy
 
-
     .. automethod:: __init__
     .. automethod:: __call__
     """
@@ -464,7 +467,6 @@ def __call__(self,
             actx: PyOpenCLArrayContext,
             source_dd: DOFDescriptorLike | None,
             dof_index: IndexList,
-            include_cluster_radii: bool = False,
             **kwargs: ArrayContainer) -> ProxyClusterGeometryData:
         """Generate proxy points for each cluster in *dof_index_set* with nodes in
         the discretization *source_dd*.
@@ -491,25 +493,14 @@ def __call__(self,
                 srcstarts=dof_index.starts)
 
         knl = self.get_radii_kernel_ex(actx)
-        _, (radii_dev,) = knl(actx.queue,
+        _, (cluster_radii,) = knl(actx.queue,
                 sources=sources,
                 srcindices=dof_index.indices,
                 srcstarts=dof_index.starts,
-                radius_factor=self.radius_factor,
                 proxy_centers=centers_dev,
                 **kwargs)
 
-        if include_cluster_radii:
-            knl = make_compute_cluster_radii_kernel_ex(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
+        radii_dev = self.radius_factor * cluster_radii
 
         # }}}
 
@@ -542,7 +533,7 @@ def __call__(self,
                 points=actx.freeze(actx.from_numpy(proxies)),
                 centers=actx.freeze(centers_dev),
                 radii=actx.freeze(radii_dev),
-                _cluster_radii=cluster_radii
+                cluster_radii=actx.freeze(cluster_radii),
                 )
 
 
@@ -564,13 +555,12 @@ def prg():
                         - sources[idim, srcindices[i + ioffset]]) ** 2)
                         ))
 
-                proxy_radius[icluster] = radius_factor * cluster_radius
+                proxy_radius[icluster] = cluster_radius
             end
             """, [
                 lp.GlobalArg("sources", None,
                     shape=(ndim, "nsources"), dim_tags="sep,C", offset=lp.auto),
                 lp.ValueArg("nsources", np.int64),
-                lp.ValueArg("radius_factor", np.float64),
                 ...
                 ],
             name="compute_cluster_radii_knl",
@@ -625,7 +615,7 @@ def prg() -> lp.ExecutorBase:
 
                 <> rqbx = rqbx_int if rqbx_ext < rqbx_int else rqbx_ext
 
-                proxy_radius[icluster] = radius_factor * rqbx
+                proxy_radius[icluster] = rqbx
             end
             """, [
                 lp.GlobalArg("sources", None,
@@ -635,7 +625,6 @@ def prg() -> lp.ExecutorBase:
                 lp.GlobalArg("center_ext", None,
                     shape=(ndim, "nsources"), dim_tags="sep,C", offset=lp.auto),
                 lp.ValueArg("nsources", np.int64),
-                lp.ValueArg("radius_factor", np.float64),
                 ...
                 ],
             name="compute_cluster_qbx_radii_knl",
@@ -668,24 +657,28 @@ def __call__(self,
             actx: PyOpenCLArrayContext,
             source_dd: DOFDescriptorLike | None,
             dof_index: IndexList,
-            include_cluster_radii: bool = False,
             **kwargs: ArrayContainer) -> ProxyClusterGeometryData:
         if source_dd is None:
             source_dd = self.places.auto_source
         dd = sym.as_dofdesc(source_dd)
 
-        radii = cast("ArrayContainer", bind(self.places, sym.expansion_radii(
-            self.ambient_dim, dofdesc=dd))(actx))
-        center_int = cast("ArrayContainer", bind(self.places, sym.expansion_centers(
-            self.ambient_dim, -1, dofdesc=dd))(actx))
-        center_ext = cast("ArrayContainer", bind(self.places, sym.expansion_centers(
-            self.ambient_dim, +1, dofdesc=dd))(actx))
+        radii = cast("ArrayContainer", bind(
+            self.places,
+            sym.expansion_radii(self.ambient_dim, dofdesc=dd)
+            )(actx))
+        center_int = cast("ArrayContainer", bind(
+            self.places,
+            sym.expansion_centers(self.ambient_dim, -1, dofdesc=dd)
+            )(actx))
+        center_ext = cast("ArrayContainer", bind(
+            self.places,
+            sym.expansion_centers(self.ambient_dim, +1, dofdesc=dd)
+            )(actx))
 
         return super().__call__(actx, dd, dof_index,
                 expansion_radii=flatten(radii, actx),
                 center_int=flatten(center_int, actx, leaf_class=DOFArray),
                 center_ext=flatten(center_ext, actx, leaf_class=DOFArray),
-                include_cluster_radii=include_cluster_radii,
                 **kwargs)
 
 # }}}

From f787cddb2a400debc0789a12701f51f617ff2c86 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 20:28:59 +0300
Subject: [PATCH 08/10] feat: allow gathering neighbors from another index list

---
 pytential/linalg/proxy.py | 71 +++++++++++++++++++++++++--------------
 1 file changed, 46 insertions(+), 25 deletions(-)

diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py
index b7729c8fb..1d9871eb0 100644
--- a/pytential/linalg/proxy.py
+++ b/pytential/linalg/proxy.py
@@ -688,24 +688,39 @@ def __call__(self,
 
 def gather_cluster_neighbor_points(
         actx: PyOpenCLArrayContext,
-        pxy: ProxyClusterGeometryData, *,
+        pxy: ProxyClusterGeometryData,
+        tgtindex: IndexList | None = None,
+        *,
         max_particles_in_box: int | None = None) -> IndexList:
-    """Generate a set of neighboring points for each cluster of points in
-    *pxy*. Neighboring points of a cluster :math:`i` are defined
-    as all the points inside the proxy ball :math:`i` that do not also
-    belong to the cluster itself.
+    r"""Generate a set of neighboring points for each cluster of points in *pxy*.
+
+    Neighboring points of a cluster :math:`i` are defined as all the points
+    from *tgtindex* that are inside the proxy ball :math:`i` but outside the
+    cluster itself. For example, given a cluster with radius :math:`r_s` and
+    proxy radius :math:`r_p > r_s`, then we gather all points such that
+    :math:`r_s < \|\mathbf{x}\| <= r_p`.
     """
 
+    srcindex = pxy.srcindex
+    if tgtindex is None:
+        tgtindex = srcindex
+
+    nclusters = srcindex.nclusters
+    if tgtindex.nclusters != nclusters:
+        raise ValueError("'tgtindex' has a different number of clusters: "
+                         f"'{tgtindex.nclusters}' (expected {nclusters})")
+
     if max_particles_in_box is None:
         max_particles_in_box = _DEFAULT_MAX_PARTICLES_IN_BOX
 
-    from pytential.source import LayerPotentialSourceBase
-
     dofdesc = pxy.dofdesc
     lpot_source = pxy.places.get_geometry(dofdesc.geometry)
-    assert isinstance(lpot_source, LayerPotentialSourceBase)
-
     discr = pxy.places.get_discretization(dofdesc.geometry, dofdesc.discr_stage)
+
+    assert (
+        dofdesc.discr_stage is None
+        or isinstance(lpot_source, QBXLayerPotentialSource)
+        ), (dofdesc, type(lpot_source))
     assert isinstance(discr, Discretization)
 
     # {{{ get only sources in the current cluster set
@@ -733,18 +748,23 @@ def prg() -> lp.ExecutorBase:
 
         return knl.executor(actx.context)
 
-    _, (sources,) = prg()(actx.queue,
+    _, (targets,) = prg()(actx.queue,
             ary=flatten(discr.nodes(), actx, leaf_class=DOFArray),
-            srcindices=pxy.srcindex.indices)
+            srcindices=tgtindex.indices)
 
     # }}}
 
     # {{{ perform area query
 
     from pytential.qbx.utils import tree_code_container
-    tcc = tree_code_container(lpot_source._setup_actx)
 
-    tree, _ = tcc.build_tree()(actx.queue, sources,
+    # NOTE: use the base source's actx for caching the code -- that has
+    # the best chance of surviving even when updating the lpot_source
+    setup_actx = discr._setup_actx
+    assert isinstance(setup_actx, PyOpenCLArrayContext)
+
+    tcc = tree_code_container(setup_actx)
+    tree, _ = tcc.build_tree()(actx.queue, targets,
             max_particles_in_box=max_particles_in_box)
     query, _ = tcc.build_area_query()(actx.queue, tree, pxy.centers, pxy.radii)
 
@@ -757,10 +777,10 @@ def prg() -> lp.ExecutorBase:
 
     pxycenters = actx.to_numpy(pxy.centers)
     pxyradii = actx.to_numpy(pxy.radii)
-    srcindex = pxy.srcindex
 
-    nbrindices: np.ndarray = np.empty(srcindex.nclusters, dtype=object)
-    for icluster in range(srcindex.nclusters):
+    eps = 100 * np.finfo(pxyradii.dtype).eps
+    nbrindices = np.empty(nclusters, dtype=object)
+    for icluster in range(nclusters):
         # get list of boxes intersecting the current ball
         istart = query.leaves_near_ball_starts[icluster]
         iend = query.leaves_near_ball_starts[icluster + 1]
@@ -779,16 +799,17 @@ def prg() -> lp.ExecutorBase:
         isources = tree.user_source_ids[isources]
 
         # get nodes inside the ball but outside the current cluster
-        # FIXME: this assumes that only the points in `pxy.secindex` should
-        # count as neighbors, not all the nodes in the discretization.
-        # FIXME: it also assumes that all the indices are sorted?
         center = pxycenters[:, icluster].reshape(-1, 1)
-        radius = pxyradii[icluster]
-        mask = ((la.norm(nodes - center, axis=0) < radius)
-                & ((isources < srcindex.starts[icluster])
-                    | (srcindex.starts[icluster + 1] <= isources)))
-
-        nbrindices[icluster] = srcindex.indices[isources[mask]]
+        radii = la.norm(nodes - center, axis=0) - eps
+        mask = (
+            (radii <= pxyradii[icluster])
+            & ((isources < tgtindex.starts[icluster])
+               | (tgtindex.starts[icluster + 1] <= isources)))
+
+        nbrindices[icluster] = tgtindex.indices[isources[mask]]
+        if nbrindices[icluster].size == 0:
+            logger.warning("Cluster '%d' has no neighbors. You might need to "
+                           "increase the proxy 'radius_factor'.", icluster)
 
     # }}}
 

From 96b91c9984d71b0c54d3d8c5d19510aeee4973a1 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Wed, 20 Aug 2025 09:18:57 +0300
Subject: [PATCH 09/10] docs: remove some unneeded references

---
 doc/conf.py               |  2 ++
 doc/misc.rst              |  9 ---------
 pytential/linalg/gmres.py | 16 ----------------
 3 files changed, 2 insertions(+), 25 deletions(-)

diff --git a/doc/conf.py b/doc/conf.py
index 03a4b5f75..f8d5e6f17 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -64,7 +64,9 @@
     "MultiVector": "obj:pymbolic.geometric_algebra.MultiVector",
     "Variable": "class:pymbolic.primitives.Variable",
     # arraycontext
+    "ArrayContainer": "obj:arraycontext.ArrayContainer",
     "ArrayOrContainerOrScalar": "obj:arraycontext.ArrayOrContainerOrScalar",
+    "ArrayOrContainerT": "obj:arraycontext.ArrayOrContainerT",
     "PyOpenCLArrayContext": "class:arraycontext.PyOpenCLArrayContext",
     "ScalarLike": "obj:arraycontext.ScalarLike",
     # modepy
diff --git a/doc/misc.rst b/doc/misc.rst
index 61be7739d..145e687c2 100644
--- a/doc/misc.rst
+++ b/doc/misc.rst
@@ -162,12 +162,3 @@ Andreas Klöckner's work on :mod:`pytential` was supported in part by
 AK also gratefully acknowledges a hardware gift from Nvidia Corporation.  The
 views and opinions expressed herein do not necessarily reflect those of the
 funding agencies.
-
-Cross-References to Other Documentation
-=======================================
-
-.. currentmodule:: numpy
-
-.. class:: int8
-
-    See :class:`numpy.generic`.
diff --git a/pytential/linalg/gmres.py b/pytential/linalg/gmres.py
index cde91bf1f..052d585a4 100644
--- a/pytential/linalg/gmres.py
+++ b/pytential/linalg/gmres.py
@@ -24,26 +24,10 @@
 """
 
 __doc__ = """
-
 .. autofunction:: gmres
-
 .. autoclass:: GMRESResult
-
 .. autoexception:: GMRESError
-
 .. autoclass:: ResidualPrinter
-
-References
-----------
-.. class:: ArrayContainer
-
-    See :class:`arraycontext.ArrayContainer`.
-
-.. currentmodule:: arraycontext.typing
-
-.. class:: ArrayOrContainerT
-
-    See :class:`arraycontext.ArrayOrContainerT`.
 """
 
 from dataclasses import dataclass

From 0c928b9981a324c6ca62d7a5dba185640ea1bd36 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Tue, 19 Aug 2025 21:37:48 +0300
Subject: [PATCH 10/10] chore: update baseline

---
 .basedpyright/baseline.json | 124 +++++++++++++++---------------------
 1 file changed, 50 insertions(+), 74 deletions(-)

diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json
index 914302583..b1b32fa5d 100644
--- a/.basedpyright/baseline.json
+++ b/.basedpyright/baseline.json
@@ -5029,7 +5029,7 @@
                 "code": "reportUnknownVariableType",
                 "range": {
                     "startColumn": 12,
-                    "endColumn": 21,
+                    "endColumn": 25,
                     "lineCount": 1
                 }
             },
@@ -5052,48 +5052,8 @@
             {
                 "code": "reportUnknownVariableType",
                 "range": {
-                    "startColumn": 12,
-                    "endColumn": 13,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownVariableType",
-                "range": {
-                    "startColumn": 16,
-                    "endColumn": 29,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAny",
-                "range": {
-                    "startColumn": 28,
-                    "endColumn": 35,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 34,
-                    "endColumn": 45,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownVariableType",
-                "range": {
-                    "startColumn": 12,
-                    "endColumn": 25,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 40,
-                    "endColumn": 53,
+                    "startColumn": 8,
+                    "endColumn": 17,
                     "lineCount": 1
                 }
             },
@@ -5204,8 +5164,16 @@
             {
                 "code": "reportUnknownArgumentType",
                 "range": {
-                    "startColumn": 31,
-                    "endColumn": 44,
+                    "startColumn": 30,
+                    "endColumn": 56,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportUnknownArgumentType",
+                "range": {
+                    "startColumn": 42,
+                    "endColumn": 55,
                     "lineCount": 1
                 }
             },
@@ -5260,25 +5228,25 @@
             {
                 "code": "reportUnknownMemberType",
                 "range": {
-                    "startColumn": 57,
-                    "endColumn": 76,
+                    "startColumn": 12,
+                    "endColumn": 31,
                     "lineCount": 1
                 }
             },
             {
                 "code": "reportUnknownArgumentType",
                 "range": {
-                    "startColumn": 62,
-                    "endColumn": 52,
-                    "lineCount": 2
+                    "startColumn": 12,
+                    "endColumn": 67,
+                    "lineCount": 1
                 }
             },
             {
                 "code": "reportUnknownArgumentType",
                 "range": {
-                    "startColumn": 62,
-                    "endColumn": 52,
-                    "lineCount": 2
+                    "startColumn": 12,
+                    "endColumn": 67,
+                    "lineCount": 1
                 }
             },
             {
@@ -5345,14 +5313,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 30,
-                    "endColumn": 53,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnknownVariableType",
                 "range": {
@@ -5485,15 +5445,23 @@
                 "code": "reportAny",
                 "range": {
                     "startColumn": 8,
-                    "endColumn": 14,
+                    "endColumn": 13,
                     "lineCount": 1
                 }
             },
             {
                 "code": "reportUnknownArgumentType",
                 "range": {
-                    "startColumn": 25,
-                    "endColumn": 39,
+                    "startColumn": 24,
+                    "endColumn": 38,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportAny",
+                "range": {
+                    "startColumn": 11,
+                    "endColumn": 36,
                     "lineCount": 1
                 }
             }
@@ -5928,16 +5896,8 @@
             {
                 "code": "reportUnknownMemberType",
                 "range": {
-                    "startColumn": 18,
-                    "endColumn": 25,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownMemberType",
-                "range": {
-                    "startColumn": 28,
-                    "endColumn": 68,
+                    "startColumn": 25,
+                    "endColumn": 65,
                     "lineCount": 1
                 }
             },
@@ -40249,6 +40209,14 @@
                     "lineCount": 4
                 }
             },
+            {
+                "code": "reportAbstractUsage",
+                "range": {
+                    "startColumn": 19,
+                    "endColumn": 18,
+                    "lineCount": 4
+                }
+            },
             {
                 "code": "reportCallIssue",
                 "range": {
@@ -40377,6 +40345,14 @@
                     "lineCount": 1
                 }
             },
+            {
+                "code": "reportAbstractUsage",
+                "range": {
+                    "startColumn": 19,
+                    "endColumn": 18,
+                    "lineCount": 4
+                }
+            },
             {
                 "code": "reportCallIssue",
                 "range": {