Port Stokes to new SystemKernels

experiments/stokes-2d-interior.py

@@ -89,7 +89,7 @@ def get_obj_array(obj_array):
     stresslet_obj = StressletWrapper(dim=2)
 
     # Describe boundary operator
-    bdry_op_sym = loc_sign * 0.5 * sigma_sym + sqrt_w * stresslet_obj.apply(inv_sqrt_w_sigma, nvec_sym, mu_sym, qbx_forced_limit='avg')
+    bdry_op_sym = loc_sign * 0.5 * sigma_sym + sqrt_w * stresslet_obj.apply(inv_sqrt_w_sigma, nvec_sym, mu_sym=mu_sym, qbx_forced_limit='avg')
 
     # Bind to the qbx discretization
     bound_op = bind(qbx, bdry_op_sym)
@@ -139,7 +139,7 @@ def couette_soln(x, y, dp, h):
     sigma = gmres_result.solution
 
     # Describe representation of solution for evaluation in domain
-    representation_sym = stresslet_obj.apply(inv_sqrt_w_sigma, nvec_sym, mu_sym, qbx_forced_limit=-2)
+    representation_sym = stresslet_obj.apply(inv_sqrt_w_sigma, nvec_sym, mu_sym=mu_sym, qbx_forced_limit=-2)
 
     from sumpy.visualization import FieldPlotter
     nsamp = 10
@@ -193,7 +193,7 @@ def stride_hack(arr):
     print("exact velocity at max error points: x -> ", err[0][max_error_loc[0]], ", y -> ", err[1][max_error_loc[1]])
 
     from pytential.symbolic.mappers import DerivativeTaker
-    rep_pressure = stresslet_obj.apply_pressure(inv_sqrt_w_sigma, nvec_sym, mu_sym, qbx_forced_limit=-2)
+    rep_pressure = stresslet_obj.apply_pressure(inv_sqrt_w_sigma, nvec_sym, mu_sym=mu_sym, qbx_forced_limit=-2)
     pressure = bind((qbx, PointsTarget(eval_points_dev)),
                      rep_pressure)(queue, sigma=sigma, mu=mu, normal=normal)
     pressure = pressure.get()
@@ -206,7 +206,7 @@ def stride_hack(arr):
     x_dir_vecs = cl.array.to_device(queue, x_dir_vecs)
     y_dir_vecs = cl.array.to_device(queue, y_dir_vecs)
     dir_vec_sym = sym.make_sym_vector("force_direction", dim)
-    rep_stress = stresslet_obj.apply_stress(inv_sqrt_w_sigma, nvec_sym, dir_vec_sym, mu_sym, qbx_forced_limit=-2)
+    rep_stress = stresslet_obj.apply_stress(inv_sqrt_w_sigma, nvec_sym, dir_vec_sym, mu_sym=mu_sym, qbx_forced_limit=-2)
 
     applied_stress_x = bind((qbx, PointsTarget(eval_points_dev)),
                              rep_stress)(queue, sigma=sigma, normal=normal, force_direction=x_dir_vecs, mu=mu)

pyproject.toml

@@ -90,15 +90,16 @@ extend-select = [
     "W",    # pycodestyle
 ]
 extend-ignore = [
-    "C90",      # McCabe complexity
+    "complex-structure",
+    "f-string",
+    "invalid-argument-name",
+    "invalid-function-name",
     "missing-whitespace-around-arithmetic-operator",
     "module-import-not-at-top-of-file",
-    "invalid-function-name",
-    "invalid-argument-name",
-    "non-lowercase-variable-in-function",
     "non-empty-init-module",
+    "non-lowercase-variable-in-function",
     "printf-string-formatting",
-    "f-string",
+    "pytest-parameter-with-default-argument",
 ]
 exclude = [
     "experiments/*.py",

pytential/symbolic/mappers.py

@@ -185,6 +185,21 @@
 
         return replace(expr, densities=densities, kernel_arguments=kernel_arguments)
 
+    def map_scalar_kernel_wrapper(
+            self, expr: pp.ScalarKernelWrapper) -> ArithmeticExpression:
+        kernel_arguments = constantdict({
+            name: componentwise(self.rec_arith, arg)
+            for name, arg in expr.kernel_arguments.items()
+        })
+        if all(arg is orig for arg, orig in zip(
+            kernel_arguments.values(),
+            expr.kernel_arguments.values(),
+            strict=True,
+        )):
+            return expr
+
+        replace(expr, kernel=expr.kernel, kernel_arguments=kernel_arguments)
+
     def map_interpolation(self, expr: pp.Interpolation):
         operand = self.rec_arith(expr.operand)
         if operand is expr.operand:
@@ -997,6 +1012,10 @@
                     expr.kernel_arguments),
                 kernel_str)
 
+    def map_scalar_kernel_wrapper(
+            self, expr: pp.ScalarKernelWrapper, enclosing_prec: int):
+        return f"{expr.kernel}"
+
     def map_interpolation(self, expr: pp.Interpolation, enclosing_prec: int):
         return "Interp[{}->{}]({})".format(
                 stringify_where(expr.from_dd),

pytential/symbolic/primitives.py

@@ -67,7 +67,7 @@
     ShapeT,
     from_numpy,
 )
-from sumpy.kernel import ScalarKernel
+from sumpy.kernel import ScalarKernel, SystemKernel
 from sumpy.symbolic import SpatialConstant
 
 from pytential.symbolic.dof_desc import (
@@ -2113,6 +2113,197 @@
             ))
 
 
+@expr_dataclass()
+class ScalarKernelWrapper(ExpressionNode):
+    kernel: ScalarKernel
+    kernel_arguments: KernelArgumentMapping
+
+    @property
+    def ambient_dim(self) -> int:
+        return self.kernel.dim
+
+    def grad(self) -> SystemKernelWrapper:
+        from sumpy.kernel import AxisSourceDerivative
+
+        result = np.empty((self.ambient_dim,), dtype=object)
+        for i in range(self.kernel.dim):
+            result[i] = ScalarKernelWrapper(
+                AxisSourceDerivative(i, self.kernel),
+                self.kernel_arguments
+            )
+
+        return SystemKernelWrapper(result)
+
+
+@expr_dataclass()
+class SystemKernelWrapper(ExpressionNode):
+    kernel: SystemKernel | None
+    components: obj_array.ObjectArrayND[ScalarKernelWrapper]
+
+    def __init__(
+            self,
+            kernel: SystemKernel | obj_array.ObjectArrayND[ScalarKernelWrapper],
+            kernel_arguments: KernelArgumentMapping | None = None,
+        ) -> None:
+        from pytools import ndindex
+
+        if kernel_arguments is None:
+            kernel_arguments = constantdict({})
+
+        if isinstance(kernel, SystemKernel):
+            knl = kernel
+            components = np.empty(kernel.shape, dtype=object)
+            for i in ndindex(kernel.shape):
+                components[i] = ScalarKernelWrapper(kernel[i], kernel_arguments)
+        else:
+            components = kernel
+            knl = None
+
+        if not components.size:
+            raise ValueError(f"system kernel cannot be empty: {components.shape}")
+
+        object.__setattr__(self, "kernel", knl)
+        object.__setattr__(self, "components", components)
+
+    @property
+    def ambient_dim(self) -> int:
+        return self.components.item(0).ambient_dim
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        return self.components.shape
+
+    @property
+    def size(self) -> int:
+        return self.components.size
+
+    @property
+    def ndim(self) -> int:
+        return self.components.ndim
+
+    # NOTE: this is required so np.einsum works with the object as is. Otherwise,
+    # it would convert it to an object array with shape `(1,)`.
+    def __array__(
+            self, dtype: Any = None, copy: bool | None = None
+    ) -> obj_array.ObjectArrayND[ScalarKernelWrapper]:
+        if not (dtype is None or np.dtype(dtype).char == "O"):
+            raise ValueError(f"non-object dtypes are not supported: {dtype}")
+
+        if copy:
+            return np.copy(self.components)
+        else:
+            return self.components
+
+    def __getitem__(self, idx: tuple[int, ...]) -> ScalarKernelWrapper:
+        return self.components[idx]
+
+    def grad(self) -> SystemKernelWrapper:
+        from pytools import ndindex
+        from sumpy.kernel import AxisSourceDerivative
+
+        result = np.empty((self.ambient_dim, *self.shape), dtype=object)
+        for i in ndindex(self.shape):
+            kernel = self.components[i:].kernel
+            kernel_arguments = self.components[1:].kernel_arguments
+
+            result[i] = ScalarKernelWrapper(
+                AxisSourceDerivative(i[0], kernel),
+                kernel_arguments
+            )
+
+        return SystemKernelWrapper(result)
+
+
+@overload
+def as_expr_kernel(
+        kernel: ScalarKernel,
+        kernel_arguments: KernelArgumentLike | None = None,
+    ) -> ScalarKernelWrapper: ...
+
+
+@overload
+def as_expr_kernel(
+        kernel: SystemKernel,
+        kernel_arguments: KernelArgumentLike | None = None,
+    ) -> SystemKernelWrapper: ...
+
+
+def as_expr_kernel(
+        kernel: ScalarKernel | SystemKernel,
+        kernel_arguments: KernelArgumentLike | None = None,
+    ) -> ScalarKernelWrapper | SystemKernelWrapper:
+    if kernel_arguments is None:
+        kernel_arguments = constantdict({})
+    else:
+        kernel_arguments = constantdict(kernel_arguments)
+
+    if isinstance(kernel, ScalarKernel):
+        return ScalarKernelWrapper(kernel, kernel_arguments)
+    elif isinstance(kernel, SystemKernel):
+        return SystemKernelWrapper(kernel, kernel_arguments)
+    else:
+        raise TypeError(f"unsupported kernel type: {type(kernel)}")
+
+
+def make_int_g(
+        expr: OperandTc,
+        *,
+        # FIXME: this should be per-kernel probably, so that we can have
+        #   L = S + D + 1/2
+        # in one go. Not clear if that's a good idea though..
+        qbx_forced_limit: QBXForcedLimit | None = None,
+        source: DOFDescriptorLike = None,
+        target: DOFDescriptorLike = None,
+    ) -> OperandTc:
+    from sumpy.kernel import SourceTransformationRemover, TargetTransformationRemover
+
+    sxr = SourceTransformationRemover()
+    txr = TargetTransformationRemover()
+
+    source = as_dofdesc(source)
+    target = as_dofdesc(target)
+
+    import pymbolic.primitives as prim
+
+    from pytential.symbolic.mappers import IdentityMapper
+
+    class IntGExpander(IdentityMapper):
+        def map_scalar_kernel_wrapper(
+                self, expr: ScalarKernelWrapper) -> ArithmeticExpression:
+            # NOTE: `map_product` should catch all of these. expressions outside
+            # of a product of the form `K * sigma` are not allowed
+            raise NotImplementedError
+
+        def map_product(self, expr: prim.Product) -> ArithmeticExpression:
+            kernel = None
+            density: ArithmeticExpression = 1
+            for child in expr.children:
+                if isinstance(child, ScalarKernelWrapper):
+                    if kernel is not None:
+                        raise ValueError(f"cannot nest kernels: {expr}")
+
+                    kernel = child
+                else:
+                    density *= self.rec(child)
+
+            if kernel is None:
+                return density
+
+            target_kernel = sxr(kernel.kernel)
+            source_kernels = (txr(kernel.kernel),)
+            return IntG(
+                target_kernel=target_kernel,
+                source_kernels=source_kernels,
+                densities=(density,),
+                qbx_forced_limit=qbx_forced_limit,
+                source=source,
+                target=target,
+                kernel_arguments=kernel.kernel_arguments,
+            )
+
+    return IntGExpander()(expr)
+
+
 _DIR_VEC_NAME = "dsource_vec"
 
 
@@ -2283,9 +2474,21 @@
              "Choosing default to '+1'.", stacklevel=2)
         qbx_forced_limit = +1
 
-    return int_g_vec(
-            kernel, density, qbx_forced_limit, source, target,
-            kernel_arguments, **kwargs)
+    if kernel_arguments is None:
+        kernel_arguments = {}
+    else:
+        kernel_arguments = dict(kernel_arguments)
+
+    if kwargs:
+        kernel_arguments = {**kernel_arguments, **kwargs}
+    kernel_arguments = constantdict(kernel_arguments)
+
+    return make_int_g(
+        as_expr_kernel(kernel, kernel_arguments) * density,
+        qbx_forced_limit=qbx_forced_limit,
+        source=source,
+        target=target,
+    )
 
 
 def tangential_derivative(