From 205cf15d57a7b782386869490180fd8b198d28ac Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 15 Mar 2026 20:55:39 +0200 Subject: [PATCH 1/2] feat: add custom __reduce__ to kernels --- sumpy/kernel.py | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/sumpy/kernel.py b/sumpy/kernel.py index b9b448a0..3861f650 100644 --- a/sumpy/kernel.py +++ b/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: From f3eac5ff4fa18c4027785b96ae20a7231ec4658e Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 15 Mar 2026 21:19:33 +0200 Subject: [PATCH 2/2] test: add test to pickle kernels --- sumpy/test/test_kernels.py | 45 +++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/sumpy/test/test_kernels.py b/sumpy/test/test_kernels.py index ae21ef03..24fa7108 100644 --- a/sumpy/test/test_kernels.py +++ b/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)'