Updated sumpy and boxtree to arraycontext branches

doc/conf.py

@@ -17,6 +17,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),

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,6 +1,6 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 from meshmode.mesh.generation import (  # noqa
         make_curve_mesh, starfish, ellipse, drop)
 

examples/helmholtz-dirichlet.py

@@ -1,7 +1,7 @@
 import numpy as np
 import numpy.linalg as la
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import \
         InterpolatoryQuadratureSimplexGroupFactory

examples/laplace-dirichlet-3d.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import \
         InterpolatoryQuadratureSimplexGroupFactory

examples/laplace-dirichlet-simple.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import \
         InterpolatoryQuadratureSimplexGroupFactory

examples/layerpot-3d.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 
 from sumpy.visualization import FieldPlotter
 from sumpy.kernel import one_kernel_2d, LaplaceKernel, HelmholtzKernel  # noqa

examples/layerpot.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 from sumpy.visualization import FieldPlotter
 from sumpy.kernel import one_kernel_2d, LaplaceKernel, HelmholtzKernel  # noqa: F401
 

examples/scaling-study.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import \
         InterpolatoryQuadratureSimplexGroupFactory

examples/cost.py → experiments/cost.py

@@ -28,7 +28,7 @@
 import numpy as np
 import pyopencl as cl
 
-from meshmode.array_context import PyOpenCLArrayContext
+from pytential.array_context import PyOpenCLArrayContext
 
 from pytential import sym, bind
 from pytential.qbx.cost import QBXCostModel
@@ -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,89 @@
+__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 meshmode.array_context import (                                # noqa: F401
+    PyOpenCLArrayContext as MeshmodePyOpenCLArrayContext)
+from sumpy.array_context import (                                   # noqa: F401
+    PyOpenCLArrayContext as SumpyPyOpenCLArrayContext,
+    make_loopy_program)
+from boxtree.array_context import dataclass_array_container         # noqa: F401
+from arraycontext.pytest import (
+        _PytestPyOpenCLArrayContextFactoryWithClass,
+        register_pytest_array_context_factory)
+
+__doc__ = """
+.. autoclass:: PyOpenCLArrayContext
+"""
+
+
+# {{{ PyOpenCLArrayContext
+
+class PyOpenCLArrayContext(SumpyPyOpenCLArrayContext):
+    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}'?")
+
+        # FIXME: this probably needs some proper logic
+        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

@@ -102,15 +102,15 @@ def partition_by_nodes(
         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)
+        tree = actx.to_numpy(tree)
         # FIXME maybe this should use IS_LEAF once available?
         leaf_boxes, = (
                 tree.box_flags & box_flags_enum.HAS_SOURCE_OR_TARGET_CHILD_BOXES == 0
@@ -707,12 +707,12 @@ def prg() -> lp.ExecutorBase:
     from pytential.qbx.utils import tree_code_container
     tcc = tree_code_container(lpot_source._setup_actx)
 
-    tree, _ = tcc.build_tree()(actx.queue, sources,
+    tree, _ = tcc.build_tree()(actx, sources,
             max_particles_in_box=max_particles_in_box)
-    query, _ = tcc.build_area_query()(actx.queue, tree, pxy.centers, pxy.radii)
+    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)
 
     # }}}