feat: add tests for system kernels

Author: alexfikl

sumpy/test/test_misc.py

@@ -59,15 +59,20 @@
 from sumpy.kernel import (
     BiharmonicKernel,
     BrinkmanletKernel,
+    BrinkmanletSystemKernel,
     BrinkmanStressKernel,
+    BrinkmanStressSystemKernel,
     ElasticityKernel,
+    ElasticitySystemKernel,
     ExpressionKernel,
     HelmholtzKernel,
     LaplaceKernel,
     LineOfCompressionKernel,
     ScalarKernel,
     StokesletKernel,
+    StokesletSystemKernel,
     StressletKernel,
+    StressletSystemKernel,
     YukawaKernel,
 )
 
@@ -778,6 +783,83 @@ def test_get_storage_index(order, knl, compressed):
 # }}}
 
 
+# {{{ test_system_kernel_components
+
+def _get_scalar_cls(knl):
+    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, cls):
+    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, cls):
+    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)'