Port test_p2e2p to use expansion toys

sumpy/test/test_kernels.py

@@ -37,7 +37,7 @@
     PyOpenCLArrayContext,
     pytest_generate_tests_for_array_contexts,
 )
-from pytools import memoize_on_first_arg, obj_array
+from pytools import memoize_on_first_arg
 from pytools.convergence import PConvergenceVerifier
 
 import sumpy.symbolic as sym
@@ -302,200 +302,114 @@ def test_p2e2p(
     res = 100
     nsources = 100
 
-    extra_kwargs = {}
+    extra_source_kwargs = {}
+    extra_kernel_kwargs = {}
     if isinstance(base_knl, HelmholtzKernel):
-        if base_knl.allow_evanescent:
-            extra_kwargs["k"] = 0.2 * (0.707 + 0.707j)
-        else:
-            extra_kwargs["k"] = 0.2
+        k = 0.2 * (0.707 + 0.707j) if base_knl.allow_evanescent else 0.2
+        extra_source_kwargs["k"] = k
+        extra_kernel_kwargs["k"] = k
     if isinstance(base_knl, StokesletKernel):
-        extra_kwargs["mu"] = 0.2
+        extra_source_kwargs["mu"] = 0.2
+        extra_kernel_kwargs["mu"] = 0.2
 
     if with_source_derivative:
         knl = DirectionalSourceDerivative(base_knl, "dir_vec")
     else:
         knl = base_knl
 
-    target_kernels = [
-            knl,
-            AxisTargetDerivative(0, knl),
-            ]
-    expn = expn_class(knl, order=order)
-
-    from sumpy import P2P, E2PFromSingleBox, P2EFromSingleBox
-    p2e = P2EFromSingleBox(expn, kernels=[knl])
-    e2p = E2PFromSingleBox(expn, kernels=target_kernels)
-    p2p = P2P(target_kernels, exclude_self=False)
-
-    from pytools.convergence import EOCRecorder
-    eoc_rec_pot = EOCRecorder()
-    eoc_rec_grad_x = EOCRecorder()
-
     from sumpy.expansion.local import LocalExpansionBase
-    if issubclass(expn_class, LocalExpansionBase):
-        h_values = [1/5, 1/7, 1/20]
-    else:
-        h_values = [1/2, 1/3, 1/5]
+    is_local = issubclass(expn_class, LocalExpansionBase)
 
     rng = np.random.default_rng(19)
     center = np.array([2, 1, 0][:knl.dim], np.float64)
-    sources = actx.from_numpy(
+    sources = (
         0.7 * (-0.5 + rng.random((knl.dim, nsources), dtype=np.float64))
         + center[:, np.newaxis])
+    strengths = np.ones(nsources, dtype=np.float64) / nsources
 
-    strengths = actx.from_numpy(np.ones(nsources, dtype=np.float64) / nsources)
-
-    source_boxes = actx.from_numpy(np.array([0], dtype=np.int32))
-    box_source_starts = actx.from_numpy(np.array([0], dtype=np.int32))
-    box_source_counts_nonchild = (
-        actx.from_numpy(np.array([nsources], dtype=np.int32)))
-
-    extra_source_kwargs = extra_kwargs.copy()
     if isinstance(knl, DirectionalSourceDerivative):
         alpha = np.linspace(0, 2*np.pi, nsources, np.float64)
         dir_vec = np.vstack([np.cos(alpha), np.sin(alpha)])
         extra_source_kwargs["dir_vec"] = actx.from_numpy(dir_vec)
 
+    if is_local:
+        toy_ctx = t.ToyContext(
+            knl,
+            local_expn_class=expn_class,
+            extra_source_kwargs=extra_source_kwargs,
+            extra_kernel_kwargs=extra_kernel_kwargs,
+        )
+    else:
+        toy_ctx = t.ToyContext(
+            knl,
+            mpole_expn_class=expn_class,
+            extra_source_kwargs=extra_source_kwargs,
+            extra_kernel_kwargs=extra_kernel_kwargs,
+        )
+
+    point_sources = t.PointSources(toy_ctx, sources, strengths)
+
+    from pytools.convergence import EOCRecorder
+    eoc_rec_pot = EOCRecorder()
+
+    h_values = [1/5, 1/7, 1/20] if is_local else [1/2, 1/3, 1/5]
+
+    rscale = 0.5  # pick something non-1
+
     from sumpy.visualization import FieldPlotter
 
+    if is_local:
+        loc_center = np.array([5.5, 0.0, 0.0][:knl.dim]) + center
+        expn = t.local_expand(actx, point_sources, loc_center,
+                              order=order, rscale=rscale)
+    else:
+        expn = t.multipole_expand(actx, point_sources, center,
+                                  order=order, rscale=rscale)
+
     for h in h_values:
-        if issubclass(expn_class, LocalExpansionBase):
-            loc_center = np.array([5.5, 0.0, 0.0][:knl.dim]) + center
-            centers = np.array(loc_center, dtype=np.float64).reshape(knl.dim, 1)
+        if is_local:
             fp = FieldPlotter(loc_center, extent=h, npoints=res)
         else:
             eval_center = np.array([1/h, 0.0, 0.0][:knl.dim]) + center
             fp = FieldPlotter(eval_center, extent=0.1, npoints=res)
-            centers = (np.array([0.0, 0.0, 0.0][:knl.dim],
-                                dtype=np.float64).reshape(knl.dim, 1)
-                        + center[:, np.newaxis])
-
-        centers = actx.from_numpy(centers)
-        targets = actx.from_numpy(obj_array.new_1d(fp.points))
 
-        rscale = 0.5  # pick something non-1
+        targets = fp.points
 
-        # {{{ apply p2e
+        pot = expn.eval(actx, targets)
+        pot_direct = point_sources.eval(actx, targets)
 
-        mpoles = p2e(actx,
-                source_boxes=source_boxes,
-                box_source_starts=box_source_starts,
-                box_source_counts_nonchild=box_source_counts_nonchild,
-                centers=centers,
-                sources=sources,
-                strengths=(strengths,),
-                nboxes=1,
-                tgt_base_ibox=0,
-                rscale=rscale,
-                **extra_source_kwargs)
-
-        # }}}
-
-        # {{{ apply e2p
-
-        ntargets = fp.points.shape[-1]
-
-        box_target_starts = actx.from_numpy(np.array([0], dtype=np.int32))
-        box_target_counts_nonchild = (
-            actx.from_numpy(np.array([ntargets], dtype=np.int32)))
-
-        pot, grad_x = e2p(
-                actx,
-                src_expansions=mpoles,
-                src_base_ibox=0,
-                target_boxes=source_boxes,
-                box_target_starts=box_target_starts,
-                box_target_counts_nonchild=box_target_counts_nonchild,
-                centers=centers,
-                targets=targets,
-                rscale=rscale,
-                **extra_kwargs)
-        pot = actx.to_numpy(pot)
-        grad_x = actx.to_numpy(grad_x)
-
-        # }}}
-
-        # {{{ compute (direct) reference solution
-
-        pot_direct, grad_x_direct = p2p(
-                actx,
-                targets, sources, (strengths,),
-                **extra_source_kwargs)
-        pot_direct = actx.to_numpy(pot_direct)
-        grad_x_direct = actx.to_numpy(grad_x_direct)
-
-        err_pot = la.norm((pot - pot_direct)/res**2)
-        err_grad_x = la.norm((grad_x - grad_x_direct)/res**2)
-
-        if 1:
-            err_pot = err_pot / la.norm((pot_direct)/res**2)
-            err_grad_x = err_grad_x / la.norm((grad_x_direct)/res**2)
-
-        if 0:
-            import matplotlib.pyplot as pt
-            from matplotlib.colors import Normalize
-
-            pt.subplot(131)
-            im = fp.show_scalar_in_matplotlib(pot.real)
-            im.set_norm(Normalize(vmin=-0.1, vmax=0.1))
-
-            pt.subplot(132)
-            im = fp.show_scalar_in_matplotlib(pot_direct.real)
-            im.set_norm(Normalize(vmin=-0.1, vmax=0.1))
-            pt.colorbar()
-
-            pt.subplot(133)
-            im = fp.show_scalar_in_matplotlib(np.log10(1e-15+np.abs(pot-pot_direct)))
-            im.set_norm(Normalize(vmin=-6, vmax=1))
-
-            pt.colorbar()
-            pt.show()
-
-        # }}}
+        err_pot = la.norm((pot - pot_direct) / res**2)
+        err_pot = err_pot / la.norm(pot_direct / res**2)
 
         eoc_rec_pot.add_data_point(h, err_pot)
-        eoc_rec_grad_x.add_data_point(h, err_grad_x)
 
     logger.info("expn_cls %s knl %s order %d", expn_class, knl, order)
     logger.info("POTENTIAL:")
     logger.info("%s", eoc_rec_pot)
-    logger.info("X TARGET DERIVATIVE:")
-    logger.info("%s", eoc_rec_grad_x)
 
     tgt_order = order + 1
-    if issubclass(expn_class, LocalExpansionBase):
-        tgt_order_grad = tgt_order - 1
+    if is_local:
         slack = 0.7
-        grad_slack = 0.5
     else:
-        tgt_order_grad = tgt_order + 1
-
         slack = 0.5
-        grad_slack = 1
 
         if order <= 2:
             slack += 1
-            grad_slack += 1
 
     if isinstance(knl, DirectionalSourceDerivative):
         slack += 1
-        grad_slack += 2
 
     if isinstance(base_knl, DirectionalSourceDerivative):
         slack += 1
-        grad_slack += 2
 
     if isinstance(base_knl, HelmholtzKernel):
         if base_knl.allow_evanescent:
             slack += 0.5
-            grad_slack += 0.5
 
         if issubclass(expn_class, VolumeTaylorMultipoleExpansionBase):
             slack += 0.3
-            grad_slack += 0.3
 
     assert eoc_rec_pot.order_estimate() > tgt_order - slack
-    assert eoc_rec_grad_x.order_estimate() > tgt_order_grad - grad_slack
 
 # }}}