port test_matrixgen to arraycontext

Author: alexfikl

test/test_matrixgen.py

@@ -20,28 +20,25 @@
 THE SOFTWARE.
 """
 
+import pytest
 import sys
+
 import numpy as np
 import numpy.linalg as la
 
-import pyopencl as cl
-import pyopencl.array  # noqa
-
-from sumpy.tools import vector_to_device
-
-import pytest
-from pyopencl.tools import (  # noqa
-        pytest_generate_tests_for_pyopencl as pytest_generate_tests)
-
+from arraycontext import pytest_generate_tests_for_array_contexts
+from sumpy.array_context import (                                 # noqa: F401
+        PytestPyOpenCLArrayContextFactory, _acf)
 
 import logging
 logger = logging.getLogger(__name__)
 
-import faulthandler
-faulthandler.enable()
+pytest_generate_tests = pytest_generate_tests_for_array_contexts([
+    PytestPyOpenCLArrayContextFactory,
+    ])
 
 
-def _build_geometry(queue, ntargets, nsources, mode, target_radius=1.0):
+def _build_geometry(actx, ntargets, nsources, mode, target_radius=1.0):
     # source points
     t = np.linspace(0.0, 2.0 * np.pi, nsources, endpoint=False)
     sources = np.array([np.cos(t), np.sin(t)])
@@ -59,14 +56,14 @@ def _build_geometry(queue, ntargets, nsources, mode, target_radius=1.0):
     centers = (1.0 - radius) * targets
     expansion_radii = np.full(ntargets, radius)
 
-    return (cl.array.to_device(queue, targets),
-            cl.array.to_device(queue, sources),
-            vector_to_device(queue, centers),
-            cl.array.to_device(queue, expansion_radii),
-            cl.array.to_device(queue, sigma))
+    return (actx.from_numpy(targets),
+            actx.from_numpy(sources),
+            actx.from_numpy(centers),
+            actx.from_numpy(expansion_radii),
+            actx.from_numpy(sigma))
 
 
-def _build_subset_indices(queue, ntargets, nsources, factor):
+def _build_subset_indices(actx, ntargets, nsources, factor):
     tgtindices = np.arange(0, ntargets)
     srcindices = np.arange(0, nsources)
 
@@ -82,17 +79,19 @@ def _build_subset_indices(queue, ntargets, nsources, factor):
 
     tgtindices, srcindices = np.meshgrid(tgtindices, srcindices)
     return (
-            cl.array.to_device(queue, tgtindices.ravel()).with_queue(None),
-            cl.array.to_device(queue, srcindices.ravel()).with_queue(None))
+            actx.freeze(actx.from_numpy(tgtindices.ravel())),
+            actx.freeze(actx.from_numpy(srcindices.ravel())))
 
 
+# {{{ test_qbx_direct
+
 @pytest.mark.parametrize("factor", [1.0, 0.6])
 @pytest.mark.parametrize("lpot_id", [1, 2])
-def test_qbx_direct(ctx_factory, factor, lpot_id):
-    logging.basicConfig(level=logging.INFO)
+def test_qbx_direct(actx_factory, factor, lpot_id, visualize=False):
+    if visualize:
+        logging.basicConfig(level=logging.INFO)
 
-    ctx = ctx_factory()
-    queue = cl.CommandQueue(ctx)
+    actx = actx_factory()
 
     ndim = 2
     order = 12
@@ -106,79 +105,88 @@ def test_qbx_direct(ctx_factory, factor, lpot_id):
         base_knl = LaplaceKernel(ndim)
         knl = DirectionalSourceDerivative(base_knl, dir_vec_name="dsource_vec")
     else:
-        raise ValueError("unknow lpot_id")
+        raise ValueError(f"unknown lpot_id: {lpot_id}")
 
     from sumpy.expansion.local import LineTaylorLocalExpansion
     expn = LineTaylorLocalExpansion(knl, order)
 
     from sumpy.qbx import LayerPotential
-    lpot = LayerPotential(ctx, expansion=expn, source_kernels=(knl,),
+    lpot = LayerPotential(actx.context, expansion=expn, source_kernels=(knl,),
             target_kernels=(base_knl,))
 
     from sumpy.qbx import LayerPotentialMatrixGenerator
-    mat_gen = LayerPotentialMatrixGenerator(ctx, expansion=expn,
-            source_kernels=(knl,), target_kernels=(base_knl,))
+    mat_gen = LayerPotentialMatrixGenerator(actx.context,
+            expansion=expn,
+            source_kernels=(knl,),
+            target_kernels=(base_knl,))
 
     from sumpy.qbx import LayerPotentialMatrixSubsetGenerator
-    blk_gen = LayerPotentialMatrixSubsetGenerator(ctx, expansion=expn,
-            source_kernels=(knl,), target_kernels=(base_knl,))
+    blk_gen = LayerPotentialMatrixSubsetGenerator(actx.context,
+            expansion=expn,
+            source_kernels=(knl,),
+            target_kernels=(base_knl,))
 
     for n in [200, 300, 400]:
         targets, sources, centers, expansion_radii, sigma = \
-                _build_geometry(queue, n, n, mode_nr, target_radius=1.2)
+                _build_geometry(actx, n, n, mode_nr, target_radius=1.2)
 
         h = 2 * np.pi / n
         strengths = (sigma * h,)
-        tgtindices, srcindices = _build_subset_indices(queue,
+        tgtindices, srcindices = _build_subset_indices(actx,
                 ntargets=n, nsources=n, factor=factor)
 
         extra_kwargs = {}
         if lpot_id == 2:
             from pytools.obj_array import make_obj_array
-            extra_kwargs["dsource_vec"] = \
-                    vector_to_device(queue, make_obj_array(np.ones((ndim, n))))
+            extra_kwargs["dsource_vec"] = (
+                    actx.from_numpy(make_obj_array(np.ones((ndim, n))))
+                    )
 
-        _, (result_lpot,) = lpot(queue,
+        _, (result_lpot,) = lpot(actx.queue,
                 targets=targets,
                 sources=sources,
                 centers=centers,
                 expansion_radii=expansion_radii,
                 strengths=strengths, **extra_kwargs)
-        result_lpot = result_lpot.get()
+        result_lpot = actx.to_numpy(result_lpot)
 
-        _, (mat,) = mat_gen(queue,
+        _, (mat,) = mat_gen(actx.queue,
                 targets=targets,
                 sources=sources,
                 centers=centers,
                 expansion_radii=expansion_radii, **extra_kwargs)
-        mat = mat.get()
-        result_mat = mat.dot(strengths[0].get())
+        mat = actx.to_numpy(mat)
+        result_mat = mat @ actx.to_numpy(strengths[0])
 
-        _, (blk,) = blk_gen(queue,
+        _, (blk,) = blk_gen(actx.queue,
                 targets=targets,
                 sources=sources,
                 centers=centers,
                 expansion_radii=expansion_radii,
                 tgtindices=tgtindices,
                 srcindices=srcindices, **extra_kwargs)
-        blk = blk.get()
+        blk = actx.to_numpy(blk)
 
-        tgtindices = tgtindices.get(queue)
-        srcindices = srcindices.get(queue)
+        tgtindices = actx.to_numpy(tgtindices)
+        srcindices = actx.to_numpy(srcindices)
 
         eps = 1.0e-10 * la.norm(result_lpot)
         assert la.norm(result_mat - result_lpot) < eps
         assert la.norm(blk - mat[tgtindices, srcindices]) < eps
 
+# }}}
+
+
+# {{{ test_p2p_direct
 
 @pytest.mark.parametrize("exclude_self", [True, False])
 @pytest.mark.parametrize("factor", [1.0, 0.6])
 @pytest.mark.parametrize("lpot_id", [1, 2])
-def test_p2p_direct(ctx_factory, exclude_self, factor, lpot_id):
-    logging.basicConfig(level=logging.INFO)
+def test_p2p_direct(actx_factory, exclude_self, factor, lpot_id, visualize=False):
+    if visualize:
+        logging.basicConfig(level=logging.INFO)
 
-    ctx = ctx_factory()
-    queue = cl.CommandQueue(ctx)
+    actx = actx_factory()
 
     ndim = 2
     mode_nr = 25
@@ -190,70 +198,73 @@ def test_p2p_direct(ctx_factory, exclude_self, factor, lpot_id):
         lknl = LaplaceKernel(ndim)
         lknl = DirectionalSourceDerivative(lknl, dir_vec_name="dsource_vec")
     else:
-        raise ValueError("unknow lpot_id")
+        raise ValueError(f"unknown lpot_id: '{lpot_id}'")
 
     from sumpy.p2p import P2P
-    lpot = P2P(ctx, [lknl], exclude_self=exclude_self)
+    lpot = P2P(actx.context, [lknl], exclude_self=exclude_self)
 
     from sumpy.p2p import P2PMatrixGenerator
-    mat_gen = P2PMatrixGenerator(ctx, [lknl], exclude_self=exclude_self)
+    mat_gen = P2PMatrixGenerator(actx.context, [lknl], exclude_self=exclude_self)
 
     from sumpy.p2p import P2PMatrixSubsetGenerator
-    blk_gen = P2PMatrixSubsetGenerator(ctx, [lknl], exclude_self=exclude_self)
+    blk_gen = P2PMatrixSubsetGenerator(
+        actx.context, [lknl], exclude_self=exclude_self)
 
     for n in [200, 300, 400]:
-        targets, sources, _, _, sigma = \
-            _build_geometry(queue, n, n, mode_nr, target_radius=1.2)
+        targets, sources, _, _, sigma = (
+            _build_geometry(actx, n, n, mode_nr, target_radius=1.2))
 
         h = 2 * np.pi / n
         strengths = (sigma * h,)
-        tgtindices, srcindices = _build_subset_indices(queue,
+        tgtindices, srcindices = _build_subset_indices(actx,
                 ntargets=n, nsources=n, factor=factor)
 
         extra_kwargs = {}
         if exclude_self:
-            extra_kwargs["target_to_source"] = \
-                cl.array.arange(queue, 0, n, dtype=np.int32)
+            extra_kwargs["target_to_source"] = (
+                actx.from_numpy(np.arange(n, dtype=np.int32))
+                )
         if lpot_id == 2:
             from pytools.obj_array import make_obj_array
-            extra_kwargs["dsource_vec"] = \
-                    vector_to_device(queue, make_obj_array(np.ones((ndim, n))))
+            extra_kwargs["dsource_vec"] = (
+                    actx.from_numpy(make_obj_array(np.ones((ndim, n)))))
 
-        _, (result_lpot,) = lpot(queue,
+        _, (result_lpot,) = lpot(actx.queue,
                 targets=targets,
                 sources=sources,
                 strength=strengths, **extra_kwargs)
-        result_lpot = result_lpot.get()
+        result_lpot = actx.to_numpy(result_lpot)
 
-        _, (mat,) = mat_gen(queue,
+        _, (mat,) = mat_gen(actx.queue,
                 targets=targets,
                 sources=sources, **extra_kwargs)
-        mat = mat.get()
-        result_mat = mat.dot(strengths[0].get())
+        mat = actx.to_numpy(mat)
+        result_mat = mat @ actx.to_numpy(strengths[0])
 
-        _, (blk,) = blk_gen(queue,
+        _, (blk,) = blk_gen(actx.queue,
                 targets=targets,
                 sources=sources,
                 tgtindices=tgtindices,
                 srcindices=srcindices, **extra_kwargs)
-        blk = blk.get()
+        blk = actx.to_numpy(blk)
 
-        tgtindices = tgtindices.get(queue)
-        srcindices = srcindices.get(queue)
+        tgtindices = actx.to_numpy(tgtindices)
+        srcindices = actx.to_numpy(srcindices)
 
         eps = 1.0e-10 * la.norm(result_lpot)
         assert la.norm(result_mat - result_lpot) < eps
         assert la.norm(blk - mat[tgtindices, srcindices]) < eps
 
+# }}}
+
 
 # You can test individual routines by typing
-# $ python test_kernels.py "test_p2p(cl.create_some_context)"
+# $ python test_matrixgen.py 'test_p2p_direct(_acf, True, 1.0, 1, visualize=True)'
 
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
     else:
-        from pytest import main
-        main([__file__])
+        pytest.main([__file__])
 
 # vim: fdm=marker