fix porting

Author: alexfikl

pytential/qbx/__init__.py

@@ -930,19 +930,21 @@ def _flat_centers(dofdesc, qbx_forced_limit):
             p2p = self.get_p2p(actx, insn.target_kernels, insn.source_kernels)
             lpot_applier_on_tgt_subset = self.get_lpot_applier_on_tgt_subset(
                     insn.target_kernels, insn.source_kernels)
+        else:
+            p2p = lpot_applier_on_tgt_subset = None
 
         for (target_name, qbx_forced_limit), outputs in other_outputs.items():
             target_discr = bound_expr.places.get_discretization(
                     target_name.geometry, target_name.discr_stage)
             flat_target_nodes = _flat_nodes(target_name)
 
             # FIXME: (Somewhat wastefully) compute P2P for all targets
-            _, output_for_each_kernel = p2p(
-                  actx,
-                  targets=flat_target_nodes,
-                  sources=flat_source_nodes,
-                  strength=flat_strengths,
-                  **flat_kernel_args)
+            assert p2p is not None
+            output_for_each_kernel = p2p(actx,
+                    targets=flat_target_nodes,
+                    sources=flat_source_nodes,
+                    strength=flat_strengths,
+                    **flat_kernel_args)
 
             target_discrs_and_qbx_sides = ((target_discr, qbx_forced_limit),)
             geo_data = self.qbx_fmm_geometry_data(
@@ -980,6 +982,7 @@ def _flat_centers(dofdesc, qbx_forced_limit):
                 tgt_subset_kwargs[f"result_{i}"] = res_i
 
             if qbx_tgt_count:
+                assert lpot_applier_on_tgt_subset is not None
                 lpot_applier_on_tgt_subset(
                         actx,
                         targets=flat_target_nodes,

pytential/qbx/fmmlib.py

@@ -24,7 +24,7 @@
 """
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -66,7 +66,7 @@ def __init__(self, actx: PyOpenCLArrayContext, *,
         # {{{ digest target_kernels
 
         ifgrad = False
-        outputs = []
+        outputs: list[tuple[Any, ...]] = []
         source_deriv_names = []
         k_names = []
 

pytential/qbx/utils.py

@@ -258,7 +258,7 @@ class TreeWithQBXMetadata(Tree):
     box_to_qbx_target_lists: Array
 
     qbx_element_to_source_starts: Array
-    qbx_element_to_center_starts: Array
+    qbx_element_to_center_starts: Array | None
 
     qbx_user_source_slice: slice
     qbx_user_center_slice: slice
@@ -399,16 +399,16 @@ def _make_centers(discr):
     del box_to_class
 
     # Compute element => source relation
-    qbx_element_to_source_starts = cast("cl_array.Array",
-                                actx.np.zeros(nelements + 1, tree.particle_id_dtype))
+    qbx_element_to_source_starts = actx.np.zeros(nelements + 1, tree.particle_id_dtype)
+
     el_offset = 0
     node_nr_base = 0
     for group in density_discr.groups:
         group_element_starts = np.arange(
                 node_nr_base, node_nr_base + group.ndofs, group.nunit_dofs,
                 dtype=tree.particle_id_dtype)
-        qbx_element_to_source_starts[el_offset:el_offset + group.nelements] = \
-                actx.from_numpy(group_element_starts)
+        qbx_element_to_source_starts[el_offset:el_offset + group.nelements] = (
+                actx.from_numpy(group_element_starts))
 
         node_nr_base += group.ndofs
         el_offset += group.nelements

pytential/unregularized.py

@@ -37,6 +37,7 @@
 from meshmode.dof_array import DOFArray
 from pytools import memoize_method
 
+from pytential.array_context import dataclass_array_container
 from pytential.source import LayerPotentialSourceBase
 
 
@@ -337,7 +338,7 @@ def copy_targets_kernel(self):
         knl = lp.tag_array_axes(knl, "targets", "stride:auto, stride:1")
         knl = lp.tag_inames(knl, {"dim": "ilp"})
 
-        return knl.executor(self.cl_context)
+        return knl.executor(self.array_context.context)
 
     @property
     @memoize_method
@@ -352,6 +353,7 @@ def build_traversal(self):
         return FMMTraversalBuilder(self.array_context)
 
 
+@dataclass_array_container
 @dataclass(frozen=True)
 class _TargetInfo:
     """