feat: add tests for system kernels

Author: alexfikl

sumpy/test/test_misc.py

@@ -59,17 +59,22 @@
 from sumpy.kernel import (
     BiharmonicKernel,
     BrinkmanletKernel,
+    BrinkmanletSystemKernel,
     BrinkmanStressKernel,
+    BrinkmanStressSystemKernel,
     ElasticityKernel,
+    ElasticitySystemKernel,
     ExpressionKernel,
+    ExpressionSystemKernel,
     HeatKernel,
     HelmholtzKernel,
-    Kernel,
     LaplaceKernel,
     LineOfCompressionKernel,
     ScalarKernel,
     StokesletKernel,
+    StokesletSystemKernel,
     StressletKernel,
+    StressletSystemKernel,
     YukawaKernel,
 )
 
@@ -88,9 +93,9 @@
 # {{{ pde check for kernels
 
 class KernelInfo:
-    kernel: Kernel
+    kernel: ScalarKernel
 
-    def __init__(self, kernel: Kernel, **kwargs):
+    def __init__(self, kernel: ScalarKernel, **kwargs: Any) -> None:
         self.kernel = kernel
         self.extra_kwargs = kwargs
         diff_op = self.kernel.get_pde_as_diff_op()
@@ -791,6 +796,83 @@ def test_get_storage_index(order, knl, compressed):
 # }}}
 
 
+# {{{ test_system_kernel_components
+
+def _get_scalar_cls(knl: ExpressionSystemKernel) -> type[ScalarKernel]:
+    if isinstance(knl, ElasticitySystemKernel):
+        return ElasticityKernel
+    elif isinstance(knl, StokesletSystemKernel):
+        return StokesletKernel
+    elif isinstance(knl, StressletSystemKernel):
+        return StressletKernel
+    elif isinstance(knl, BrinkmanletSystemKernel):
+        return BrinkmanletKernel
+    elif isinstance(knl, BrinkmanStressSystemKernel):
+        return BrinkmanStressKernel
+    else:
+        raise AssertionError
+
+
+@pytest.mark.parametrize("dim", [2, 3])
+@pytest.mark.parametrize("cls", [
+    ElasticitySystemKernel,
+    StokesletSystemKernel,
+    StressletSystemKernel,
+    BrinkmanletSystemKernel,
+    BrinkmanStressSystemKernel,
+])
+def test_system_kernel_components(dim: int, cls: type[ExpressionSystemKernel]) -> None:
+    knl = cls(dim=dim)
+    scalar_cls = _get_scalar_cls(knl)
+    shape = knl.shape
+
+    from itertools import product
+
+    if len(shape) == 2:
+        assert shape == (dim, dim)
+        for i, j in product(range(dim), repeat=2):
+            comp = knl[i, j]
+            expected_ij = tuple(sorted([i, j]))
+            assert isinstance(comp, scalar_cls)
+            assert comp.dim == dim
+            assert (comp.icomp, comp.jcomp) == expected_ij
+            # symmetry
+            assert knl[i, j] is knl[j, i]
+    else:
+        assert shape == (dim, dim, dim)
+        for i, j, k in product(range(dim), repeat=3):
+            comp = knl[i, j, k]
+            expected_ijk = tuple(sorted([i, j, k]))
+
+            assert isinstance(comp, scalar_cls)
+            assert comp.dim == dim
+            assert (comp.icomp, comp.jcomp, comp.kcomp) == expected_ijk
+            # symmetry: all permutations of sorted indices are the same
+            s = tuple(sorted([i, j, k]))
+            assert knl[i, j, k] is knl[s]
+
+# }}}
+
+
+# {{{ test_system_kernel_pickle
+
+@pytest.mark.parametrize("dim", [2, 3])
+@pytest.mark.parametrize("cls", [
+    ElasticitySystemKernel,
+    StokesletSystemKernel,
+    StressletSystemKernel,
+    BrinkmanletSystemKernel,
+    BrinkmanStressSystemKernel,
+])
+def test_system_kernel_pickle(dim: int, cls: type[ExpressionSystemKernel]) -> None:
+    from pickle import dumps, loads
+
+    knl = cls(dim=dim)
+    assert loads(dumps(knl)) == knl
+
+# }}}
+
+
 # You can test individual routines by typing
 # $ python test_misc.py 'test_pde_check_kernels(_acf,
 #       KernelInfo(HelmholtzKernel(2), k=5), order=5)'