Updated sumpy and boxtree to arraycontext branches

doc/conf.py

@@ -18,6 +18,10 @@
         "DOFDescriptorLike": "pytential.symbolic.dof_desc.DOFDescriptorLike",
         }
 
+nitpick_ignore_regex = [
+        ["py:class", r"boxtree\.array_context\.(.+)"],
+        ]
+
 intersphinx_mapping = {
     "arraycontext": ("https://documen.tician.de/arraycontext", None),
     "boxtree": ("https://documen.tician.de/boxtree", None),
@@ -100,11 +104,13 @@
     "P2P": "class:sumpy.p2p.P2P",
     "P2PBase": "class:sumpy.p2p.P2PBase",
     "FMMLevelToOrder": "class:sumpy.fmm.FMMLevelToOrder",
+    "MultipoleExpansionFromOrderFactory": "class:sumpy.fmm.MultipoleExpansionFromOrderFactory",  # noqa: E501
+    "LocalExpansionFromOrderFactory": "class:sumpy.fmm.LocalExpansionFromOrderFactory",
     # pytential
-    "ExpressionNode": "class:pytential.symbolic.primitives.ExpressionNode",
     "DOFDescriptorLike": "data:pytential.symbolic.dof_desc.DOFDescriptorLike",
     "DOFGranularity": "data:pytential.symbolic.dof_desc.DOFGranularity",
     "DiscretizationStage": "data:pytential.symbolic.dof_desc.DiscretizationStage",
+    "ExpressionNode": "class:pytential.symbolic.primitives.ExpressionNode",
     "GeometryId": "data:pytential.symbolic.dof_desc.GeometryId",
     "KernelArgumentLike": "obj:pytential.symbolic.primitives.KernelArgumentLike",
     "KernelArgumentMapping": "obj:pytential.symbolic.primitives.KernelArgumentMapping",
@@ -114,10 +120,8 @@
     "Side": "obj:pytential.symbolic.primitives.Side",
     "TargetOrDiscretization": "obj:pytential.target.TargetOrDiscretization",
     "VectorExpression": "obj:pytential.symbolic.pde.scalar.VectorExpression",
-    "pytential.symbolic.dof_desc.DOFDescriptorLike":
-        "data:pytential.symbolic.dof_desc.DOFDescriptorLike",
-    "pytential.symbolic.primitives.ExpressionNode":
-        "class:pytential.symbolic.primitives.ExpressionNode",
+    "pytential.symbolic.dof_desc.DOFDescriptorLike": "data:pytential.symbolic.dof_desc.DOFDescriptorLike",  # noqa: E501
+    "pytential.symbolic.primitives.ExpressionNode": "class:pytential.symbolic.primitives.ExpressionNode",  # noqa: E501
     "sym.DOFDescriptor": "class:pytential.symbolic.dof_desc.DOFDescriptor",
     "sym.IntG": "class:pytential.symbolic.primitives.IntG",
     "sym.var": "obj:pytential.symbolic.primitives.var",

doc/qbx.rst

@@ -8,6 +8,11 @@ QBX internals
     This page documents :mod:`pytential` internals and is not typically
     needed in end-user applications.
 
+Array Context
+-------------
+
+.. automodule:: pytential.array_context
+
 Refinement
 ----------
 

examples/fmm-error.py

@@ -1,10 +1,11 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.mesh.generation import drop, ellipse, make_curve_mesh, starfish  # noqa
 from sumpy.kernel import HelmholtzKernel, LaplaceKernel
 from sumpy.visualization import FieldPlotter
 
+from pytential.array_context import PyOpenCLArrayContext
+
 
 def main():
     import logging

examples/helmholtz-dirichlet.py

@@ -1,13 +1,13 @@
 import numpy as np
 import numpy.linalg as la
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 
@@ -55,7 +55,7 @@
         meshes = [
                 affine_map(
                     base_mesh,
                     A=np.diag([dx*0.25, dx*0.25]),
                     b=np.array([dx*(i_x-nx/2), dx*(i_y-ny/2)]))
                 for i_x in range(nx)
                 for i_y in range(ny)]

examples/laplace-dirichlet-3d.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 
@@ -48,7 +48,7 @@
         meshes = [
                 affine_map(
                     base_mesh,
                     A=np.diag([1, 1, 1]),
                     b=np.array([0, 0, iz*dz]))
                 for iz in range(nz)]
 

examples/laplace-dirichlet-simple.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 

examples/layerpot-3d.py

@@ -1,10 +1,10 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from sumpy.kernel import HelmholtzKernel, LaplaceKernel, one_kernel_2d  # noqa
 from sumpy.visualization import FieldPlotter
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 
 
 target_order = 5

examples/layerpot.py

@@ -1,11 +1,11 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.mesh.generation import drop, ellipse, starfish  # noqa: F401
 from sumpy.kernel import HelmholtzKernel, LaplaceKernel, one_kernel_2d  # noqa: F401
 from sumpy.visualization import FieldPlotter
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 
 
 target_order = 16

examples/scaling-study.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import (
     InterpolatoryQuadratureSimplexGroupFactory,
 )
 
 from pytential import bind, sym
+from pytential.array_context import PyOpenCLArrayContext
 from pytential.target import PointsTarget
 
 
@@ -38,7 +38,7 @@
     meshes = [
             affine_map(
                 base_mesh,
                 A=np.diag([dx*0.25, dx*0.25]),
                 b=np.array([dx*(i_x-nx/2), dx*(i_y-ny/2)]))
             for i_x in range(nx)
             for i_y in range(ny)]

examples/cost.py → experiments/cost.py

@@ -138,8 +138,7 @@ def calibrate_cost_model(ctx):
 
         for _ in range(RUNS):
             timing_data = {}
-            bound_op.eval({"sigma": sigma}, array_context=actx,
-                    timing_data=timing_data)
+            bound_op.eval({"sigma": sigma}, array_context=actx)
 
             model_results.append(modeled_cost)
             timing_results.append(timing_data)
@@ -175,8 +174,7 @@ def test_cost_model(ctx, calibration_params):
         temp_timing_results = []
         for _ in range(RUNS):
             timing_data = {}
-            bound_op.eval({"sigma": sigma},
-                    array_context=actx, timing_data=timing_data)
+            bound_op.eval({"sigma": sigma}, array_context=actx)
             temp_timing_results.append(one(timing_data.values()))
 
         timing_result = {}

pytential/array_context.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+
+__copyright__ = "Copyright (C) 2022 Alexandru Fikl"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+from arraycontext.pytest import (
+    _PytestPyOpenCLArrayContextFactoryWithClass,
+    register_pytest_array_context_factory,
+)
+from sumpy.array_context import (  # noqa: F401
+    PyOpenCLArrayContext as PyOpenCLArrayContextBase,
+    make_loopy_program,
+)
+
+
+__doc__ = """
+.. autoclass:: PyOpenCLArrayContext
+"""
+
+
+# {{{ PyOpenCLArrayContext
+
+class PyOpenCLArrayContext(PyOpenCLArrayContextBase):
+    def transform_loopy_program(self, t_unit):
+        kernel = t_unit.default_entrypoint
+        options = kernel.options
+
+        if not options.return_dict or not options.no_numpy:
+            raise ValueError(
+                "loopy kernels passed to 'call_loopy' must have 'return_dict' "
+                "and 'no_numpy' options set. Did you use 'make_loopy_program' "
+                f"to create the kernel '{kernel.name}'?")
+
+        # apply transformations from sumpy
+        t_unit = super().transform_loopy_program(t_unit)
+
+        # apply transformations from meshmode
+        from meshmode.array_context import _transform_loopy_inner
+        transformed_t_unit = _transform_loopy_inner(t_unit)
+
+        if transformed_t_unit is not None:
+            return transformed_t_unit
+
+        return t_unit
+
+# }}}
+
+
+# {{{ pytest
+
+def _acf():
+    import pyopencl as cl
+    ctx = cl.create_some_context()
+    queue = cl.CommandQueue(ctx)
+
+    return PyOpenCLArrayContext(queue, force_device_scalars=True)
+
+
+class PytestPyOpenCLArrayContextFactory(
+        _PytestPyOpenCLArrayContextFactoryWithClass):
+    actx_class = PyOpenCLArrayContext
+
+    def __call__(self):
+        # NOTE: prevent any cache explosions during testing!
+        from sympy.core.cache import clear_cache
+        clear_cache()
+
+        return super().__call__()
+
+
+register_pytest_array_context_factory(
+    "pytential.pyopencl",
+    PytestPyOpenCLArrayContextFactory)
+
+# }}}

pytential/linalg/proxy.py

@@ -120,28 +120,29 @@
         from pytential.qbx.utils import tree_code_container
         tcc = tree_code_container(lpot_source._setup_actx)
 
-        tree, _ = tcc.build_tree()(actx.queue,
+        tree, _ = tcc.build_tree()(actx,
                 particles=flatten(
                     actx.thaw(discr.nodes()), actx, leaf_class=DOFArray
                     ),
                 max_particles_in_box=max_particles_in_box,
                 kind=tree_kind)
 
         from boxtree import box_flags_enum
-        tree = tree.get(actx.queue)
-        # FIXME: maybe this should use IS_LEAF once available?
+
+        tree = actx.to_numpy(tree)
+        # FIXME maybe this should use IS_LEAF once available?
         assert tree.box_flags is not None
         leaf_boxes, = (
                 tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0
                 ).nonzero()
 
         indices = np.empty(len(leaf_boxes), dtype=object)
         starts: onp.Array1D[np.integer] | None = None
 
         for i, ibox in enumerate(leaf_boxes):
             box_start = tree.box_source_starts[ibox]
             box_end = box_start + tree.box_source_counts_cumul[ibox]
             indices[i] = tree.user_source_ids[box_start:box_end]
     else:
         if discr.ambient_dim != 2 and discr.dim == 1:
             raise ValueError("only curves are supported for 'tree_kind=None'")
@@ -766,12 +767,11 @@
     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)
+    tree, _ = tcc.build_tree()(actx, targets, max_particles_in_box=max_particles_in_box)
+    query, _ = tcc.build_area_query()(actx, tree, pxy.centers, pxy.radii)
 
-    tree = tree.get(actx.queue)
-    query = query.get(actx.queue)
+    tree = actx.to_numpy(tree)
+    query = actx.to_numpy(query)
 
     # }}}
 
@@ -784,9 +784,9 @@
     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]
         iboxes = query.leaves_near_ball_lists[istart:iend]
 
         if (iend - istart) <= 0:
             nbrindices[icluster] = np.empty(0, dtype=np.int64)