Add custom __reduce__ to kernels

sumpy/kernel.py

@@ -466,6 +466,10 @@ def __init__(self, dim: int):
             global_scaling_const=1,
         )
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (self.__class__, (self.dim,))
+
     @property
     @override
     def is_complex_valued(self) -> bool:
@@ -496,6 +500,10 @@ def __init__(self, dim: int) -> None:
 
         super().__init__(dim, expression=expr, global_scaling_const=scaling)
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (self.__class__, (self.dim,))
+
     @property
     @override
     def is_complex_valued(self) -> bool:
@@ -545,6 +553,10 @@ def __init__(self, dim: int) -> None:
 
         super().__init__(dim, expression=expr, global_scaling_const=scaling)
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (self.__class__, (self.dim,))
+
     @property
     @override
     def is_complex_valued(self) -> bool:
@@ -611,6 +623,13 @@ def __init__(self,
         object.__setattr__(self, "helmholtz_k_name", helmholtz_k_name)
         object.__setattr__(self, "allow_evanescent", allow_evanescent)
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (
+            self.__class__,
+            (self.dim, self.helmholtz_k_name, self.allow_evanescent),
+        )
+
     @property
     @override
     def is_complex_valued(self) -> bool:
@@ -698,6 +717,10 @@ def __init__(self, dim: int, yukawa_lambda_name: str = "lam") -> None:
         super().__init__(dim, expression=expr, global_scaling_const=scaling)
         object.__setattr__(self, "yukawa_lambda_name", yukawa_lambda_name)
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (self.__class__, (self.dim, self.yukawa_lambda_name))
+
     @property
     @override
     def is_complex_valued(self) -> bool:
@@ -964,6 +987,13 @@ def __init__(self,
         object.__setattr__(self, "kcomp", kcomp)
         object.__setattr__(self, "viscosity_mu", mu)
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (
+            self.__class__,
+            (self.dim, self.icomp, self.jcomp, self.kcomp, self.viscosity_mu),
+        )
+
     @property
     @override
     def is_complex_valued(self) -> bool:
@@ -1052,6 +1082,13 @@ def __init__(self,
         object.__setattr__(self, "viscosity_mu", mu)
         object.__setattr__(self, "poisson_ratio", nu)
 
+    @override
+    def __reduce__(self) -> tuple[object, ...]:
+        return (
+            self.__class__,
+            (self.dim, self.axis, self.viscosity_mu, self.poisson_ratio),
+        )
+
     @property
     @override
     def is_complex_valued(self) -> bool:

sumpy/test/test_kernels.py

@@ -26,6 +26,7 @@
 import logging
 import sys
 from functools import partial
+from typing import TYPE_CHECKING
 
 import numpy as np
 import numpy.linalg as la
@@ -36,7 +37,7 @@
     PyOpenCLArrayContext,
     pytest_generate_tests_for_array_contexts,
 )
-from pytools import obj_array
+from pytools import memoize_on_first_arg, obj_array
 from pytools.convergence import PConvergenceVerifier
 
 import sumpy.symbolic as sym
@@ -62,15 +63,22 @@
     AxisTargetDerivative,
     BiharmonicKernel,
     DirectionalSourceDerivative,
+    ElasticityKernel,
     HelmholtzKernel,
     Kernel,
     LaplaceKernel,
+    LineOfCompressionKernel,
+    OneKernel,
     StokesletKernel,
+    StressletKernel,
     YukawaKernel,
 )
 from sumpy.test.geometries import make_ellipsoid, make_torus
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 logger = logging.getLogger(__name__)
 
 pytest_generate_tests = pytest_generate_tests_for_array_contexts([
@@ -857,6 +865,8 @@ def test_m2m_compressed_error_helmholtz(actx_factory: ArrayContextFactory, dim,
 # }}}
 
 
+# {{{ test_jump
+
 @pytest.mark.parametrize(("kernel_cls", "kernel_kwargs"), [
                          (LaplaceKernel, {}),
                          (HelmholtzKernel, {"k": 1}),
@@ -922,6 +932,39 @@ def test_jump(
     err = abs((inside-outside) - -1)
     assert err < tol, err
 
+# }}}
+
+
+# {{{ test_pickle
+
+@memoize_on_first_arg
+def get_kernel_name_for_test(knl: Kernel) -> Callable[[str], str]:
+    return lambda prefix: f"{prefix}: {type(knl).__name__}"
+
+
+@pytest.mark.parametrize("knl", [
+    BiharmonicKernel(2),
+    ElasticityKernel(2, 0, 0),
+    HelmholtzKernel(3, helmholtz_k_name="kay"),
+    LaplaceKernel(3),
+    LineOfCompressionKernel(),
+    OneKernel(2),
+    StokesletKernel(2, 0, 0),
+    StressletKernel(2, 0, 0, 0),
+    YukawaKernel(2, yukawa_lambda_name="lambda"),
+])
+def test_pickle(knl: Kernel) -> None:
+    import pickle
+
+    result = pickle.dumps(knl)
+    assert pickle.loads(result) == knl
+
+    _ = get_kernel_name_for_test(knl)
+    result = pickle.dumps(knl)
+    assert pickle.loads(result) == knl
+
+# }}}
+
 
 # You can test individual routines by typing
 # $ python test_kernels.py 'test_p2p(_acf, True)'