diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 97c79771..a2eaef37 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -9104,18 +9104,18 @@ } }, { - "code": "reportAny", + "code": "reportUnknownMemberType", "range": { - "startColumn": 12, - "endColumn": 17, + "startColumn": 15, + "endColumn": 24, "lineCount": 1 } }, { - "code": "reportUnknownMemberType", + "code": "reportAny", "range": { - "startColumn": 19, - "endColumn": 28, + "startColumn": 12, + "endColumn": 17, "lineCount": 1 } }, @@ -9159,6 +9159,14 @@ "lineCount": 1 } }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 21, + "lineCount": 1 + } + }, { "code": "reportArgumentType", "range": { @@ -9223,6 +9231,198 @@ "lineCount": 1 } }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 51, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 36, + "endColumn": 60, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 49, + "endColumn": 59, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 46, + "endColumn": 83, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 46, + "endColumn": 83, + "lineCount": 1 + } + }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 78, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 22, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 36, + "endColumn": 60, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 49, + "endColumn": 59, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 46, + "endColumn": 83, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 46, + "endColumn": 83, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 36, + "endColumn": 60, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 49, + "endColumn": 59, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 46, + "endColumn": 86, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 46, + "endColumn": 86, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 22, + "endColumn": 66, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 36, + "endColumn": 60, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 49, + "endColumn": 59, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 46, + "endColumn": 83, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 46, + "endColumn": 83, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 36, + "endColumn": 60, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 49, + "endColumn": 59, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 46, + "endColumn": 86, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 46, + "endColumn": 86, + "lineCount": 1 + } + }, { "code": "reportIncompatibleMethodOverride", "range": { @@ -21092,26 +21292,10 @@ } }, { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 41, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 41, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", + "code": "reportAny", "range": { - "startColumn": 8, - "endColumn": 25, + "startColumn": 47, + "endColumn": 53, "lineCount": 1 } }, @@ -21180,10 +21364,10 @@ } }, { - "code": "reportUnknownMemberType", + "code": "reportAny", "range": { - "startColumn": 50, - "endColumn": 67, + "startColumn": 45, + "endColumn": 46, "lineCount": 1 } }, @@ -22610,6 +22794,102 @@ "endColumn": 56, "lineCount": 1 } + }, + { + "code": "reportCallIssue", + "range": { + "startColumn": 10, + "endColumn": 22, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 20, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 25, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 32, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 37, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 20, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 25, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 32, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 37, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 44, + "endColumn": 54, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 49, + "endColumn": 54, + "lineCount": 1 + } + }, + { + "code": "reportCallIssue", + "range": { + "startColumn": 10, + "endColumn": 22, + "lineCount": 1 + } } ], "./sumpy/test/test_qbx.py": [ diff --git a/sumpy/expansion/__init__.py b/sumpy/expansion/__init__.py index eee9f882..a79ce51f 100644 --- a/sumpy/expansion/__init__.py +++ b/sumpy/expansion/__init__.py @@ -48,7 +48,7 @@ from sumpy.expansion.diff_op import MultiIndex from sumpy.expansion.local import LocalExpansionBase from sumpy.expansion.multipole import MultipoleExpansionBase - from sumpy.kernel import Kernel, KernelArgument + from sumpy.kernel import KernelArgument, ScalarKernel logger = logging.getLogger(__name__) @@ -99,7 +99,7 @@ class ExpansionBase(ABC): .. automethod:: __ne__ """ - kernel: Kernel + kernel: ScalarKernel order: int use_rscale: bool = field(kw_only=True, default=True) @@ -152,7 +152,7 @@ def get_coefficient_identifiers(self) -> Sequence[MultiIndex]: @abstractmethod def coefficients_from_source(self, - kernel: Kernel, + kernel: ScalarKernel, avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -171,7 +171,7 @@ def coefficients_from_source(self, """ def coefficients_from_source_vec(self, - kernels: Sequence[Kernel], + kernels: Sequence[ScalarKernel], avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -198,7 +198,7 @@ def coefficients_from_source_vec(self, return result def loopy_expansion_formation(self, - kernels: Sequence[Kernel], + kernels: Sequence[ScalarKernel], strength_usage: Sequence[int], nstrengths: int ) -> lp.TranslationUnit: @@ -212,7 +212,7 @@ def loopy_expansion_formation(self, @abstractmethod def evaluate(self, - kernel: Kernel, + kernel: ScalarKernel, coeffs: Sequence[sym.Expr], bvec: sym.Matrix, rscale: sym.Expr, @@ -224,7 +224,7 @@ def evaluate(self, in *coeffs*. """ - def loopy_evaluator(self, kernels: Sequence[Kernel]) -> lp.TranslationUnit: + def loopy_evaluator(self, kernels: Sequence[ScalarKernel]) -> lp.TranslationUnit: """ :returns: a :mod:`loopy` kernel that returns the evaluated target transforms of the potential given by *kernels*. @@ -236,7 +236,7 @@ def loopy_evaluator(self, kernels: Sequence[Kernel]) -> lp.TranslationUnit: # {{{ copy - def with_kernel(self, kernel: Kernel) -> ExpansionBase: + def with_kernel(self, kernel: ScalarKernel) -> ExpansionBase: return replace(self, kernel=kernel) def copy(self, **kwargs: Any) -> Self: @@ -566,7 +566,7 @@ class LinearPDEBasedExpansionTermsWrangler(ExpansionTermsWrangler): .. automethod:: __init__ """ - knl: Kernel + knl: ScalarKernel @override def get_coefficient_identifiers(self) -> Sequence[MultiIndex]: @@ -973,7 +973,7 @@ class MultipoleExpansionFactory(Protocol): .. automethod:: __call__ """ def __call__(self, - kernel: Kernel, + kernel: ScalarKernel, order: int, *, use_rscale: bool = True ) -> MultipoleExpansionBase: @@ -987,7 +987,7 @@ class LocalExpansionFactory(Protocol): .. automethod:: __call__ """ def __call__(self, - kernel: Kernel, + kernel: ScalarKernel, order: int, *, use_rscale: bool = True ) -> LocalExpansionBase: @@ -1002,7 +1002,7 @@ class ExpansionFactoryBase(ABC): @abstractmethod def get_local_expansion_class(self, - base_kernel: Kernel, / + base_kernel: ScalarKernel, / ) -> LocalExpansionFactory: """ :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*. @@ -1010,7 +1010,7 @@ def get_local_expansion_class(self, @abstractmethod def get_multipole_expansion_class(self, - base_kernel: Kernel, / + base_kernel: ScalarKernel, / ) -> MultipoleExpansionFactory: """ :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*. @@ -1024,7 +1024,7 @@ class VolumeTaylorExpansionFactory(ExpansionFactoryBase): @override def get_local_expansion_class( - self, base_kernel: Kernel, / + self, base_kernel: ScalarKernel, / ) -> type[LocalExpansionBase]: """ :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*. @@ -1034,7 +1034,7 @@ def get_local_expansion_class( @override def get_multipole_expansion_class( - self, base_kernel: Kernel, / + self, base_kernel: ScalarKernel, / ) -> type[MultipoleExpansionBase]: """ :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*. @@ -1050,7 +1050,7 @@ class DefaultExpansionFactory(ExpansionFactoryBase): @override def get_local_expansion_class(self, - base_kernel: Kernel, /, + base_kernel: ScalarKernel, /, ) -> LocalExpansionFactory: """ :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*. @@ -1067,7 +1067,7 @@ def get_local_expansion_class(self, @override def get_multipole_expansion_class(self, - base_kernel: Kernel, /, + base_kernel: ScalarKernel, /, ) -> MultipoleExpansionFactory: """ :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*. diff --git a/sumpy/expansion/level_to_order.py b/sumpy/expansion/level_to_order.py index 0c975da6..791e2f67 100644 --- a/sumpy/expansion/level_to_order.py +++ b/sumpy/expansion/level_to_order.py @@ -43,7 +43,7 @@ from collections.abc import Sequence import sumpy.symbolic as sym - from sumpy.kernel import Kernel + from sumpy.kernel import ScalarKernel class TreeLike(Protocol): @@ -71,7 +71,7 @@ class FMMLibExpansionOrderFinder: """ def __call__(self, - kernel: Kernel, + kernel: ScalarKernel, kernel_args: dict[str, sym.Expr] | Sequence[tuple[str, sym.Expr]], tree: TreeLike, level: int) -> int: @@ -163,7 +163,7 @@ class SimpleExpansionOrderFinder: """ def __call__(self, - kernel: Kernel, + kernel: ScalarKernel, kernel_args: dict[str, sym.Expr] | Sequence[tuple[str, sym.Expr]], tree: TreeLike, level: int) -> int: diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py index afeaaceb..86e56167 100644 --- a/sumpy/expansion/local.py +++ b/sumpy/expansion/local.py @@ -56,7 +56,7 @@ HankelBased2DMultipoleExpansion, MultipoleExpansionBase, ) - from sumpy.kernel import Kernel + from sumpy.kernel import ScalarKernel logger = logging.getLogger(__name__) @@ -131,7 +131,7 @@ def get_coefficient_identifiers(self) -> Sequence[MultiIndex]: @override def coefficients_from_source(self, - kernel: Kernel, + kernel: ScalarKernel, avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -171,7 +171,7 @@ def coefficients_from_source(self, @override def evaluate(self, - kernel: Kernel, + kernel: ScalarKernel, coeffs: Sequence[sym.Expr], bvec: sym.Matrix, rscale: sym.Expr, @@ -224,7 +224,7 @@ def m2l_translation(self) -> M2LTranslationBase: @override def coefficients_from_source_vec(self, - kernels: Sequence[Kernel], + kernels: Sequence[ScalarKernel], avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -272,7 +272,7 @@ def save_temp(x: sym.Expr) -> sym.Expr: @override def coefficients_from_source(self, - kernel: Kernel, + kernel: ScalarKernel, avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -283,7 +283,7 @@ def coefficients_from_source(self, @override def evaluate(self, - kernel: Kernel, + kernel: ScalarKernel, coeffs: Sequence[sym.Expr], bvec: sym.Matrix, rscale: sym.Expr, @@ -556,7 +556,7 @@ def get_coefficient_identifiers(self) -> Sequence[MultiIndex]: @override def coefficients_from_source(self, - kernel: Kernel, + kernel: ScalarKernel, avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -580,7 +580,7 @@ def coefficients_from_source(self, @override def evaluate(self, - kernel: Kernel, + kernel: ScalarKernel, coeffs: Sequence[sym.Expr], bvec: sym.Matrix, rscale: sym.Expr, diff --git a/sumpy/expansion/loopy.py b/sumpy/expansion/loopy.py index a52cfbb5..9944183c 100644 --- a/sumpy/expansion/loopy.py +++ b/sumpy/expansion/loopy.py @@ -42,14 +42,16 @@ from pymbolic.typing import ArithmeticExpression from sumpy.expansion import ExpansionBase - from sumpy.kernel import Kernel + from sumpy.kernel import ScalarKernel logger = logging.getLogger(__name__) def make_e2p_loopy_kernel( - expansion: ExpansionBase, kernels: Sequence[Kernel]) -> lp.TranslationUnit: + expansion: ExpansionBase, + kernels: Sequence[ScalarKernel], + ) -> lp.TranslationUnit: """ A helper function that creates a :mod:`loopy` kernel for multipole/local evaluation. @@ -152,7 +154,7 @@ def make_e2p_loopy_kernel( def make_p2e_loopy_kernel( expansion: ExpansionBase, - kernels: Sequence[Kernel], + kernels: Sequence[ScalarKernel], strength_usage: Sequence[int], nstrengths: int) -> lp.TranslationUnit: """ diff --git a/sumpy/expansion/m2l.py b/sumpy/expansion/m2l.py index 5a77df51..92b43f73 100644 --- a/sumpy/expansion/m2l.py +++ b/sumpy/expansion/m2l.py @@ -59,7 +59,7 @@ from sumpy.assignment_collection import SymbolicAssignmentCollection from sumpy.expansion.diff_op import MultiIndex - from sumpy.kernel import Kernel + from sumpy.kernel import ScalarKernel logger = logging.getLogger(__name__) @@ -88,7 +88,7 @@ class M2LTranslationClassFactoryBase(ABC): @abstractmethod def get_m2l_translation_class( self, - base_kernel: Kernel, + base_kernel: ScalarKernel, local_expansion_class: type[LocalExpansionBase] ) -> type[M2LTranslationBase]: """ @@ -105,7 +105,7 @@ class NonFFTM2LTranslationClassFactory(M2LTranslationClassFactoryBase): @override def get_m2l_translation_class( self, - base_kernel: Kernel, + base_kernel: ScalarKernel, local_expansion_class: type[LocalExpansionBase] ) -> type[M2LTranslationBase]: from sumpy.expansion.local import ( @@ -129,7 +129,7 @@ class FFTM2LTranslationClassFactory(M2LTranslationClassFactoryBase): @override def get_m2l_translation_class( self, - base_kernel: Kernel, + base_kernel: ScalarKernel, local_expansion_class: type[LocalExpansionBase] ) -> type[M2LTranslationBase]: from sumpy.expansion.local import ( @@ -153,7 +153,7 @@ class DefaultM2LTranslationClassFactory(M2LTranslationClassFactoryBase): @override def get_m2l_translation_class( self, - base_kernel: Kernel, + base_kernel: ScalarKernel, local_expansion_class: type[LocalExpansionBase] ) -> type[M2LTranslationBase]: from sumpy.expansion.local import ( diff --git a/sumpy/expansion/multipole.py b/sumpy/expansion/multipole.py index 197e076e..a9f98ce7 100644 --- a/sumpy/expansion/multipole.py +++ b/sumpy/expansion/multipole.py @@ -45,7 +45,7 @@ from sumpy.assignment_collection import SymbolicAssignmentCollection from sumpy.expansion.diff_op import MultiIndex - from sumpy.kernel import Kernel + from sumpy.kernel import ScalarKernel logger = logging.getLogger(__name__) @@ -76,7 +76,7 @@ class VolumeTaylorMultipoleExpansionBase(VolumeTaylorExpansionMixin, @override def coefficients_from_source_vec(self, - kernels: Sequence[Kernel], + kernels: Sequence[ScalarKernel], avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -116,7 +116,7 @@ def coefficients_from_source_vec(self, @override def coefficients_from_source( self, - kernel: Kernel, + kernel: ScalarKernel, avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -130,7 +130,7 @@ def coefficients_from_source( @override def evaluate(self, - kernel: Kernel, + kernel: ScalarKernel, coeffs: Sequence[sym.Expr], bvec: sym.Matrix, rscale: sym.Expr, @@ -440,7 +440,7 @@ def get_coefficient_identifiers(self) -> Sequence[MultiIndex]: @override def coefficients_from_source( self, - kernel: Kernel, + kernel: ScalarKernel, avec: sym.Matrix, bvec: sym.Matrix | None, rscale: sym.Expr, @@ -469,7 +469,7 @@ def coefficients_from_source( @override def evaluate(self, - kernel: Kernel, + kernel: ScalarKernel, coeffs: Sequence[sym.Expr], bvec: sym.Matrix, rscale: sym.Expr, diff --git a/sumpy/fmm.py b/sumpy/fmm.py index c53c40b8..12149303 100644 --- a/sumpy/fmm.py +++ b/sumpy/fmm.py @@ -61,7 +61,7 @@ P2EFromSingleBox, P2PFromCSR, ) -from sumpy.kernel import Kernel +from sumpy.kernel import ScalarKernel from sumpy.tools import ( get_opencl_fft_app, run_opencl_fft, @@ -105,8 +105,8 @@ class SumpyTreeIndependentDataForWrangler(TreeIndependentDataForWrangler): multipole_expansion_factory: MultipoleExpansionFromOrderFactory local_expansion_factory: LocalExpansionFromOrderFactory - source_kernels: Sequence[Kernel] | None - target_kernels: Sequence[Kernel] + source_kernels: Sequence[ScalarKernel] | None + target_kernels: Sequence[ScalarKernel] exclude_self: bool use_rscale: bool | None strength_usage: Sequence[int] | None @@ -115,11 +115,11 @@ def __init__(self, array_context: ArrayContext, multipole_expansion_factory: MultipoleExpansionFromOrderFactory, local_expansion_factory: LocalExpansionFromOrderFactory, - target_kernels: Sequence[Kernel], + target_kernels: Sequence[ScalarKernel], exclude_self: bool = False, use_rscale: bool | None = None, strength_usage: Sequence[int] | None = None, - source_kernels: Sequence[Kernel] | None = None): + source_kernels: Sequence[ScalarKernel] | None = None): """ :arg multipole_expansion_factory: a callable of a single argument (order) that returns a multipole expansion. @@ -257,7 +257,7 @@ def app() -> Any: # {{{ expansion wrangler FMMLevelToOrder: TypeAlias = Callable[ - [Kernel, frozenset[tuple[str, object]], Tree, int], + [ScalarKernel, frozenset[tuple[str, object]], Tree, int], int] diff --git a/sumpy/kernel.py b/sumpy/kernel.py index e0a7ba79..1ea437ad 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -37,6 +37,7 @@ ) import numpy as np +from constantdict import constantdict from typing_extensions import Self, override import loopy as lp @@ -44,7 +45,7 @@ from pymbolic import Expression, var from pymbolic.mapper import CSECachingMapperMixin, IdentityMapper from pymbolic.primitives import make_sym_vector -from pytools import memoize_method +from pytools import memoize_method, obj_array import sumpy.symbolic as sym from sumpy.derivative_taker import ( @@ -53,11 +54,10 @@ ExprDerivativeTaker, diff_derivative_coeff_dict, ) -from sumpy.symbolic import SpatialConstant, pymbolic_real_norm_2 if TYPE_CHECKING: - from collections.abc import Callable, Iterable, Sequence + from collections.abc import Callable, Iterable, Mapping, Sequence import sympy as sp @@ -69,7 +69,11 @@ ---------------- .. autoclass:: KernelArgument -.. autoclass:: Kernel +.. autoclass:: GenericKernel + :show-inheritance: +.. autoclass:: ScalarKernel + :show-inheritance: +.. autoclass:: SystemKernel :show-inheritance: Symbolic kernels @@ -79,33 +83,47 @@ :show-inheritance: :members: mapper_method -PDE kernels ------------ +.. autoclass:: ExpressionSystemKernel + :show-inheritance: + :members: mapper_method + +.. autoclass:: OneKernel + :show-inheritance: + :members: mapper_method + +Scalar PDE kernels +------------------ .. autoclass:: LaplaceKernel :show-inheritance: :members: mapper_method + .. autoclass:: BiharmonicKernel :show-inheritance: :members: mapper_method + .. autoclass:: HelmholtzKernel :show-inheritance: :members: mapper_method + .. autoclass:: YukawaKernel :show-inheritance: :members: mapper_method + .. autoclass:: StokesletKernel :show-inheritance: :members: mapper_method .. autoclass:: StressletKernel :show-inheritance: :members: mapper_method + .. autoclass:: ElasticityKernel :show-inheritance: :members: mapper_method .. autoclass:: LineOfCompressionKernel :show-inheritance: :members: mapper_method + .. autoclass:: BrinkmanletKernel :show-inheritance: :members: mapper_method @@ -116,6 +134,33 @@ :show-inheritance: :members: mapper_method +System PDE Kernels +------------------ + +.. autoclass:: ElasticitySystemKernel + :show-inheritance: + :members: mapper_method + +.. autoclass:: StokesSystemKernel + :show-inheritance: + :members: mapper_method +.. autoclass:: StokesletSystemKernel + :show-inheritance: + :members: mapper_method +.. autoclass:: StressletSystemKernel + :show-inheritance: + :members: mapper_method + +.. autoclass:: BrinkmanSystemKernel + :show-inheritance: + :members: mapper_method +.. autoclass:: BrinkmanletSystemKernel + :show-inheritance: + :members: mapper_method +.. autoclass:: BrinkmanStressSystemKernel + :show-inheritance: + :members: mapper_method + .. [Pozrikidis1992] C. Pozrikidis, *Boundary Integral and Singularity Methods for Linearized Viscous Flow*, Cambridge University Press, 1992. @@ -198,8 +243,25 @@ def name(self) -> str: # {{{ basic kernel interface + +def _diff(expr: sym.Expr, vec: sp.Matrix, mi: tuple[int, ...]) -> sym.Expr: + """Take the derivative of an expression.""" + dim = len(mi) + assert vec.shape == (dim, 1) + + for i in range(dim): + if mi[i] == 0: + continue + expr = expr.diff(vec[i], mi[i]) + + return expr + + +KernelObjectT = TypeVar("KernelObjectT") + + @dataclass(frozen=True) -class Kernel(ABC): +class GenericKernel(ABC): """Basic kernel interface. .. autoattribute:: mapper_method @@ -208,32 +270,19 @@ class Kernel(ABC): .. autoattribute:: dim .. autoproperty:: is_complex_valued - .. automethod:: get_base_kernel - .. automethod:: replace_base_kernel - .. automethod:: prepare_loopy_kernel - .. automethod:: get_code_transformer - .. automethod:: get_expression - .. automethod:: postprocess_at_source - .. automethod:: postprocess_at_target .. automethod:: get_global_scaling_const - .. automethod:: get_args - .. automethod:: get_source_args + .. automethod:: get_expression .. automethod:: get_pde_as_diff_op """ - dim: int - """Dimension of the space the kernel is defined in.""" - # TODO: Allow kernels that are not translation invariant is_translation_invariant: ClassVar[bool] = True """A boolean flag indicating whether the kernel is translation invariant.""" mapper_method: ClassVar[str] """The name of the mapper method called for the kernel.""" - @property - @abstractmethod - def is_complex_valued(self) -> bool: - """A boolean flag indicating whether this kernel is complex valued.""" + dim: int + """Dimension of the space the kernel is defined in.""" @override def __repr__(self) -> str: @@ -249,19 +298,82 @@ def __repr__(self) -> str: return f"{type(self).__name__}({', '.join(args)})" - def get_base_kernel(self) -> Kernel: + @property + @abstractmethod + def is_complex_valued(self) -> bool: + """A boolean flag indicating whether this kernel is complex valued.""" + + @abstractmethod + def get_expression(self, dist_vec: sym.Matrix) -> sym.Expr: + """ + :returns: a :mod:`sympy` expression for the kernel. + """ + + @abstractmethod + def get_global_scaling_const(self) -> sym.Expr: + r"""A global scaling constant of the kernel. + + Typically, this ensures that the kernel is scaled so that + :math:`\mathcal{L}(G)(x) = C \delta(x)` with a constant of 1, where + :math:`\mathcal{L}` is the PDE operator associated with the kernel. Not + to be confused with *rscale*, which keeps expansion coefficients + benignly scaled. + """ + + @abstractmethod + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: + r""" + :returns: the PDE for the kernel as a + :class:`sumpy.expansion.diff_op.LinearPDESystemOperator` object + :math:`\mathcal{L}`, where :math:`\mathcal{L}(u) = 0` is the PDE. + """ + + +@dataclass(frozen=True, repr=False) +class ScalarKernel(GenericKernel, ABC): + """Scalar kernel interface. + + .. automethod:: get_base_kernel + .. automethod:: replace_base_kernel + + .. automethod:: get_args + .. automethod:: get_source_args + .. automethod:: prepare_loopy_kernel + .. automethod:: get_code_transformer + + .. automethod:: get_derivative_taker + .. automethod:: get_derivative_coeff_dict_at_source + .. automethod:: postprocess_at_source + .. automethod:: postprocess_at_target + """ + + def get_base_kernel(self) -> ScalarKernel: """ :returns: the kernel being wrapped by this one, or else *self*. """ return self - def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel: + def replace_base_kernel(self, new_base_kernel: ScalarKernel) -> ScalarKernel: """ :returns: the base kernel being wrapped by this one, or else *new_base_kernel*. """ return new_base_kernel + def get_args(self) -> Sequence[KernelArgument]: + """ + :returns: list of :class:`KernelArgument` instances describing extra + arguments used by the kernel. + """ + return [] + + def get_source_args(self) -> Sequence[KernelArgument]: + """ + :returns: list of :class:`KernelArgument` instances describing extra + arguments used by kernel in picking up contributions from point sources. + """ + return [] + def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit: """Apply some changes (such as registering function manglers) to the kernel. @@ -272,37 +384,11 @@ def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationU def get_code_transformer(self) -> Callable[[Expression], Expression]: """ :returns: a function to postprocess the :mod:`pymbolic` expression - generated from the result of :meth:`get_expression` on the way to - code generation. + generated from the result of :meth:`~GenericKernel.get_expression` + on the way to code generation. """ return lambda expr: expr - @abstractmethod - def get_expression(self, dist_vec: sym.Matrix) -> sym.Expr: - """ - :returns: a :mod:`sympy` expression for the kernel. - """ - - @abstractmethod - def get_pde_as_diff_op(self) -> LinearPDESystemOperator: - r""" - :returns: the PDE for the kernel as a - :class:`sumpy.expansion.diff_op.LinearPDESystemOperator` object - :math:`\mathcal{L}`, where :math:`\mathcal{L}(u) = 0` is the PDE. - """ - - def _diff(self, - expr: sym.Expr, - vec: sp.Matrix, - mi: tuple[int, ...]) -> sym.Expr: - """Take the derivative of an expression.""" - for i in range(self.dim): - if mi[i] == 0: - continue - expr = expr.diff(vec[i], mi[i]) - - return expr - @abstractmethod def get_derivative_taker( self, @@ -316,6 +402,19 @@ def get_derivative_taker( *dvec*. """ + def get_derivative_coeff_dict_at_source( + self, expr_dict: DerivativeCoeffDict, + ) -> DerivativeCoeffDict: + r"""Get the derivative transformation of the expression at the source. + + The transformation is represented by the *expr_dict* which maps from a + multi-index *mi* to a coefficient *coeff*. The Expression represented by + *expr_dict* is :math:`\sum_{mi} \frac{\partial^mi}{x^mi}G * coeff`. + + This function is meant to be overridden by child classes where necessary. + """ + return expr_dict + @overload def postprocess_at_source( self, expr: sym.Expr, avec: sym.Matrix @@ -343,7 +442,7 @@ def postprocess_at_source( result = 0 for mi, coeff in expr_dict.items(): - result += coeff * self._diff(expr, avec, mi) + result += coeff * _diff(expr, avec, mi) assert isinstance(result, sym.Expr) return result @@ -373,49 +472,50 @@ def postprocess_at_target(self, """ return expr - def get_derivative_coeff_dict_at_source( - self, expr_dict: DerivativeCoeffDict, - ) -> DerivativeCoeffDict: - r"""Get the derivative transformation of the expression at the source. - The transformation is represented by the *expr_dict* which maps from a - multi-index *mi* to a coefficient *coeff*. The Expression represented by - *expr_dict* is :math:`\sum_{mi} \frac{\partial^mi}{x^mi}G * coeff`. +@dataclass(frozen=True, repr=False) +class Kernel(ScalarKernel, ABC): + """Deprecated kernel interface.""" - This function is meant to be overridden by child classes where necessary. - """ - return expr_dict + def __post_init__(self) -> None: + from warnings import warn - @abstractmethod - def get_global_scaling_const(self) -> sym.Expr: - r"""A global scaling constant of the kernel. + warn("'sumpy.kernel.Kernel' is deprecated and will be removed in Q2 2027. " + "Use 'ScalarKernel' or 'SystemKernel' instead, as appropriate.", + DeprecationWarning, stacklevel=3) - Typically, this ensures that the kernel is scaled so that - :math:`\mathcal{L}(G)(x) = C \delta(x)` with a constant of 1, where - :math:`\mathcal{L}` is the PDE operator associated with the kernel. Not - to be confused with *rscale*, which keeps expansion coefficients - benignly scaled. - """ - def get_args(self) -> Sequence[KernelArgument]: - """ - :returns: list of :class:`KernelArgument` instances describing extra - arguments used by the kernel. - """ - return [] +@dataclass(frozen=True, repr=False) +class SystemKernel(GenericKernel, ABC): + @property + def ndim(self) -> int: + return len(self.shape) - def get_source_args(self) -> Sequence[KernelArgument]: + @property + @abstractmethod + def shape(self) -> tuple[int, ...]: + pass + + @abstractmethod + def __getitem__(self, index: tuple[int, ...], /) -> ScalarKernel: + pass + + @abstractmethod + @override + def get_expression( # pyright: ignore[reportIncompatibleMethodOverride] + self, dist_vec: sym.Matrix, / + ) -> obj_array.ObjectArrayND[sym.Expr]: """ - :returns: list of :class:`KernelArgument` instances describing extra - arguments used by kernel in picking up contributions from point sources. + :returns: a :mod:`sympy` expression for the kernel. """ - return [] # }}} +# {{{ generic expression kernel + @dataclass(frozen=True, repr=False) -class ExpressionKernel(Kernel, ABC): +class ExpressionKernel(ScalarKernel, ABC): r""" .. autoattribute:: expression .. autoattribute:: global_scaling_const @@ -446,24 +546,21 @@ def __str__(self) -> str: @override def get_expression(self, dist_vec: sym.Matrix) -> sym.Expr: - from sumpy.symbolic import PymbolicToSympyMapperWithSymbols - expr = PymbolicToSympyMapperWithSymbols().to_expr(self.expression) + expr = sym.PymbolicToSympyMapperWithSymbols().to_expr(self.expression) if self.dim != len(dist_vec): raise ValueError( "'dist_vec' length does not match expected dimension: " f"kernel dim is '{self.dim}' and dist_vec has length '{len(dist_vec)}'") - from sumpy.symbolic import Symbol return expr.xreplace({ - Symbol(f"d{i}"): dist_vec_i + sym.Symbol(f"d{i}"): dist_vec_i for i, dist_vec_i in enumerate(dist_vec) }) @override def get_global_scaling_const(self) -> sym.Expr: - from sumpy.symbolic import PymbolicToSympyMapperWithSymbols - return PymbolicToSympyMapperWithSymbols().to_expr(self.global_scaling_const) + return sym.PymbolicToSympyMapperWithSymbols().to_expr(self.global_scaling_const) @override def get_derivative_taker( @@ -501,6 +598,47 @@ def is_complex_valued(self) -> bool: one_kernel_3d = OneKernel(3) +@dataclass(frozen=True, repr=False) +class ExpressionSystemKernel(SystemKernel, ABC): + mapper_method: ClassVar[str] = "map_expression_system_kernel" + + components: Mapping[tuple[int, ...], ScalarKernel] + """A multi-dimensional array of scalar kernels.""" + global_scaling_const: Expression + """See :attr:`sumpy.kernel.ExpressionKernel.global_scaling_const`.""" + + @override + def __str__(self) -> str: + return f"VecExprKnl{self.dim}D" + + @override + def __getitem__(self, index: tuple[int, ...], /) -> ScalarKernel: + return self.components[index] + + @override + def get_expression( + self, dist_vec: sym.Matrix, / + ) -> obj_array.ObjectArrayND[sym.Expr]: + from pytools import ndindex + + result = np.empty(self.shape, dtype=object) + for i in ndindex(result.shape): + result[i] = self.components[i].get_expression(dist_vec) + + return result + + @override + def get_global_scaling_const(self) -> sym.Expr: + return sym.PymbolicToSympyMapperWithSymbols().to_expr(self.global_scaling_const) + + @override + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: + raise NotImplementedError + + +# }}} + + # {{{ PDE kernels class LaplaceKernel(ExpressionKernel): @@ -516,11 +654,11 @@ class LaplaceKernel(ExpressionKernel): def __init__(self, dim: int) -> None: # See (Kress LIE, Thm 6.2) for scaling if dim == 2: - r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + r = sym.pymbolic_real_norm_2(make_sym_vector("d", dim)) expr = var("log")(r) scaling = 1/(-2*var("pi")) elif dim == 3: - r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + r = sym.pymbolic_real_norm_2(make_sym_vector("d", dim)) expr = 1/r scaling = 1/(4*var("pi")) else: @@ -569,7 +707,7 @@ class BiharmonicKernel(ExpressionKernel): mapper_method: ClassVar[str] = "map_biharmonic_kernel" def __init__(self, dim: int) -> None: - r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + r = sym.pymbolic_real_norm_2(make_sym_vector("d", dim)) if dim == 2: # Ref: Farkas, Peter. Mathematical foundations for fast algorithms # for the biharmonic equation. Technical Report 765, @@ -643,17 +781,17 @@ def __init__(self, dim: int, helmholtz_k_name: str = "k", allow_evanescent: bool = False) -> None: - k = SpatialConstant(helmholtz_k_name) + k = sym.SpatialConstant(helmholtz_k_name) # Guard against code using the old positional interface. assert isinstance(allow_evanescent, bool) if dim == 2: - r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + r = sym.pymbolic_real_norm_2(make_sym_vector("d", dim)) expr = var("hankel_1")(0, k*r) scaling = var("I")/4 elif dim == 3: - r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + r = sym.pymbolic_real_norm_2(make_sym_vector("d", dim)) expr = var("exp")(var("I")*k*r)/r scaling = 1/(4*var("pi")) else: @@ -689,9 +827,8 @@ def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationU def get_args(self) -> Sequence[KernelArgument]: k_dtype = np.complex128 if self.allow_evanescent else np.float64 return [ - KernelArgument( - loopy_arg=lp.ValueArg(self.helmholtz_k_name, k_dtype), - )] + KernelArgument(loopy_arg=lp.ValueArg(self.helmholtz_k_name, k_dtype)), + ] @override def get_derivative_taker( @@ -732,7 +869,7 @@ class YukawaKernel(ExpressionKernel): """ def __init__(self, dim: int, yukawa_lambda_name: str = "lam") -> None: - lam = SpatialConstant(yukawa_lambda_name) + lam = sym.SpatialConstant(yukawa_lambda_name) # NOTE: The Yukawa kernel is given by [1] # -1/(2 pi)**(n/2) * (lam/r)**(n/2-1) * K(n/2-1, lam r) @@ -743,7 +880,7 @@ def __init__(self, dim: int, yukawa_lambda_name: str = "lam") -> None: # [3] https://dlmf.nist.gov/10.47#E2 # [4] https://dlmf.nist.gov/10.49 - r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + r = sym.pymbolic_real_norm_2(make_sym_vector("d", dim)) if dim == 2: # NOTE: transform K(0, lam r) into a Hankel function using [2] expr = var("hankel_1")(0, var("I")*lam*r) @@ -786,9 +923,8 @@ def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationU @override def get_args(self) -> Sequence[KernelArgument]: return [ - KernelArgument( - loopy_arg=lp.ValueArg(self.yukawa_lambda_name, np.float64), - )] + KernelArgument(loopy_arg=lp.ValueArg(self.yukawa_lambda_name, np.float64)), + ] @override def get_derivative_taker( @@ -858,11 +994,11 @@ def __init__(self, f"'poisson_ratio_name' is not a str: {type(poisson_ratio_name)}" ) - mu = SpatialConstant(viscosity_mu_name) - nu = SpatialConstant(poisson_ratio_name) + mu = sym.SpatialConstant(viscosity_mu_name) + nu = sym.SpatialConstant(poisson_ratio_name) d = make_sym_vector("d", dim) - r = pymbolic_real_norm_2(d) + r = sym.pymbolic_real_norm_2(d) delta_ij = 1 if icomp == jcomp else 0 if dim == 2: @@ -961,10 +1097,10 @@ def __init__(self, raise TypeError( f"'viscosity_mu_name' is not a str: {type(viscosity_mu_name)}" ) - mu = SpatialConstant(viscosity_mu_name) + mu = sym.SpatialConstant(viscosity_mu_name) d = make_sym_vector("d", dim) - r = pymbolic_real_norm_2(d) + r = sym.pymbolic_real_norm_2(d) delta_ij = 1 if icomp == jcomp else 0 if dim == 2: @@ -1060,7 +1196,7 @@ def __init__(self, ) d = make_sym_vector("d", dim) - r = pymbolic_real_norm_2(d) + r = sym.pymbolic_real_norm_2(d) if dim == 2: expr = d[icomp]*d[jcomp]*d[kcomp]/r**4 @@ -1160,12 +1296,12 @@ def __init__(self, f"'poisson_ratio_name' is not a str: {type(poisson_ratio_name)}" ) - mu = SpatialConstant(viscosity_mu_name) - nu = SpatialConstant(poisson_ratio_name) + mu = sym.SpatialConstant(viscosity_mu_name) + nu = sym.SpatialConstant(poisson_ratio_name) if dim == 3: d = make_sym_vector("d", dim) - r = pymbolic_real_norm_2(d) + r = sym.pymbolic_real_norm_2(d) # Kelvin solution expr = d[axis] * var("log")(r + d[axis]) - r @@ -1252,11 +1388,11 @@ def __init__(self, jcomp: int, viscosity_mu_name: str = "mu", darcy_impermeability_name: str = "k") -> None: - mu = SpatialConstant(viscosity_mu_name) - k = SpatialConstant(darcy_impermeability_name) + mu = sym.SpatialConstant(viscosity_mu_name) + k = sym.SpatialConstant(darcy_impermeability_name) d = make_sym_vector("d", dim) - r = pymbolic_real_norm_2(d) + r = sym.pymbolic_real_norm_2(d) R = k * r # noqa: N806 delta_ij = 1 if icomp == jcomp else 0 @@ -1333,6 +1469,7 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: w = make_identity_diff_op(self.dim) k = sym.Symbol(self.darcy_impermeability_name) + return laplacian(laplacian(w) - k**2 * w) @@ -1377,10 +1514,10 @@ def __init__(self, kcomp: int, viscosity_mu_name: str = "mu", darcy_impermeability_name: str = "k") -> None: - k = SpatialConstant(darcy_impermeability_name) + k = sym.SpatialConstant(darcy_impermeability_name) d = make_sym_vector("d", dim) - r = pymbolic_real_norm_2(d) + r = sym.pymbolic_real_norm_2(d) R = k * r # noqa: N806 delta_ij = 1 if icomp == jcomp else 0 delta_ik = 1 if icomp == kcomp else 0 @@ -1475,61 +1612,411 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: @dataclass(frozen=True, repr=False) class HeatKernel(ExpressionKernel): - r"""The Green's function for the heat equation given by - :math:`e^{-r^2/{4 \alpha t}}/\sqrt{(4 \pi \alpha t)^d}` - where :math:`d` is the number of spatial dimensions. + r"""The Green's function for the heat equation. + + .. math:: + + \frac{\partial}{\partial t} K(t, \mathbf{x}, \mathbf{y}) + - \alpha \Delta K(t, \mathbf{x}, \mathbf{y}) + = \delta(t) \delta(\mathbf{x} - \mathbf{y}) .. note:: - This kernel cannot be used in an FMM yet and can only - be used in expansions and evaluations that occur forward - in the time dimension. + This kernel cannot be used in an FMM yet and can only be used in + expansions and evaluations that occur forward in the time dimension. """ - heat_alpha_name: str mapper_method: ClassVar[str] = "map_heat_kernel" + heat_alpha_name: str + def __init__(self, spatial_dims: int, heat_alpha_name: str = "alpha"): dim = spatial_dims + 1 + alpha = sym.SpatialConstant(heat_alpha_name) + d = make_sym_vector("d", dim) t = d[-1] - r = pymbolic_real_norm_2(d[:-1]) - alpha = SpatialConstant(heat_alpha_name) + r = sym.pymbolic_real_norm_2(d[:-1]) + expr = var("exp")(-r**2/(4 * alpha * t)) / var("sqrt")(t**(dim - 1)) scaling = 1/var("sqrt")((4*var("pi")*alpha)**(dim - 1)) - super().__init__( - dim, - expression=expr, - global_scaling_const=scaling, - ) - + super().__init__(dim, expression=expr, global_scaling_const=scaling) object.__setattr__(self, "heat_alpha_name", heat_alpha_name) + @override + def __reduce__(self) -> tuple[object, ...]: + return (self.__class__, (self.dim - 1, self.heat_alpha_name)) + @property @override def is_complex_valued(self) -> bool: return False @override - def __repr__(self): + def __str__(self): return f"HeatKnl{self.dim - 1}D" @override def get_args(self): return [ - KernelArgument( - loopy_arg=lp.ValueArg(self.heat_alpha_name, np.float64), - )] + KernelArgument(loopy_arg=lp.ValueArg(self.heat_alpha_name, np.float64)) + ] @override - def get_pde_as_diff_op(self): + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: from sumpy.expansion.diff_op import diff, laplacian, make_identity_diff_op + alpha = sym.Symbol(self.heat_alpha_name) w = make_identity_diff_op(self.dim - 1, time_dependent=True) - time_diff = [0]*self.dim - time_diff[-1] = 1 - return diff(w, tuple(time_diff)) - laplacian(w) * alpha + t_mi = (*([0] * (self.dim - 1)), 1) + + return diff(w, t_mi) - alpha * laplacian(w) + + +# }}} + + +# {{{ vector PDE kernels + + +@dataclass(frozen=True, repr=False) +class ElasticitySystemKernel(ExpressionSystemKernel): + r"""A kernel for the linear elasticity (Navier-Cauchy) equations + (see e.g. Section 2.2 in [Hsiao2008]_). + + This kernel uses :class:`ElasticityKernel` for its components. + """ + + mapper_method: ClassVar[str] = "map_elasticity_system_kernel" + + viscosity_mu_name: str + r"""The argument name to use for the dynamic viscosity :math:`\mu` when + generating functions to evaluate this kernel. + """ + poisson_ratio_name: str + r"""The argument name to use for Poisson's ratio :math:`\nu` when generating + functions to evaluate this kernel. + """ + + def __init__(self, + dim: int, + viscosity_mu_name: str = "mu", + poisson_ratio_name: str = "nu") -> None: + components = {} + + for i in range(dim): + for j in range(i, dim): + components[i, j] = components[j, i] = ElasticityKernel( + dim, i, j, + viscosity_mu_name=viscosity_mu_name, + poisson_ratio_name=poisson_ratio_name, + ) + + super().__init__(dim=dim, + components=constantdict(components), + global_scaling_const=components[0, 0].global_scaling_const) + object.__setattr__(self, "viscosity_mu_name", viscosity_mu_name) + object.__setattr__(self, "poisson_ratio_name", poisson_ratio_name) + + @override + def __str__(self) -> str: + return ( + f"ElasticityKnl{self.dim}D" + f"({self.viscosity_mu_name}, {self.poisson_ratio_name})") + + @override + def __reduce__(self): + return ( + type(self), + (self.dim, self.viscosity_mu_name, self.poisson_ratio_name)) + + @property + @override + def is_complex_valued(self) -> bool: + return False + + @property + @override + def shape(self) -> tuple[int, ...]: + return (self.dim, self.dim) + + @override + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: + from sumpy.expansion.diff_op import ( + divergence, + gradient, + laplacian, + make_identity_diff_op, + ) + + mu = sym.Symbol(self.viscosity_mu_name) + nu = sym.Symbol(self.poisson_ratio_name) + u = make_identity_diff_op(self.dim, self.dim) + + return mu * laplacian(u) + mu / (1 - 2 * nu) * gradient(divergence(u)) + + +@dataclass(frozen=True, repr=False) +class StokesSystemKernel(ExpressionSystemKernel, ABC): + """A kernel for the Stokes equations.""" + + mapper_method: ClassVar[str] = "map_stokes_system_kernel" + + viscosity_mu_name: str + r"""The argument name to use for the dynamic viscosity :math:`\mu` when + generating functions to evaluate this kernel. + """ + + @property + @override + def is_complex_valued(self) -> bool: + return False + + @override + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: + from sumpy.expansion.diff_op import ( + concat, + divergence, + gradient, + laplacian, + make_identity_diff_op, + ) + + mu = sym.Symbol(self.viscosity_mu_name) + u_and_p = make_identity_diff_op(self.dim, self.dim + 1) + u = u_and_p[:self.dim] + p = u_and_p[self.dim] + + return concat(mu * laplacian(u) - gradient(p), divergence(u)) + + +@dataclass(frozen=True, repr=False) +class StokesletSystemKernel(StokesSystemKernel): + r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_). + + This kernel uses :class:`StokesletKernel` for its components. + """ + + mapper_method: ClassVar[str] = "map_stokeslet_system_kernel" + + def __init__(self, + dim: int, + viscosity_mu_name: str = "mu") -> None: + components = {} + + for i in range(dim): + for j in range(i, dim): + components[i, j] = components[j, i] = StokesletKernel( + dim, i, j, + viscosity_mu_name=viscosity_mu_name, + ) + + super().__init__(dim=dim, + components=constantdict(components), + global_scaling_const=components[0, 0].global_scaling_const, + viscosity_mu_name=viscosity_mu_name) + + @override + def __str__(self) -> str: + return ( + f"StokesletKnl{self.dim}D({self.viscosity_mu_name})") + + @override + def __reduce__(self): + return (type(self), (self.dim, self.viscosity_mu_name)) + + @property + @override + def shape(self) -> tuple[int, ...]: + return (self.dim, self.dim) + + +@dataclass(frozen=True, repr=False) +class StressletSystemKernel(StokesSystemKernel): + r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_). + + This kernel uses :class:`StressletKernel` for its components. + """ + + mapper_method: ClassVar[str] = "map_stresslet_system_kernel" + + def __init__(self, + dim: int, + viscosity_mu_name: str = "mu") -> None: + components = {} + + from itertools import product + for i, j, k in product(range(dim), repeat=3): + if j >= i and k >= j: + components[i, j, k] = StressletKernel( + dim, i, j, k, + viscosity_mu_name=viscosity_mu_name, + ) + else: + # NOTE: this exploits the symmetry of the Stresslet -- that + # :math:`T_{012}` is identical to :math:`T_{120}` -- and avoids + # creating multiple kernels with the same expression. + s = cast("tuple[int, int, int]", tuple(sorted([i, j, k]))) + components[i, j, k] = components[s] + + super().__init__(dim=dim, + components=constantdict(components), + global_scaling_const=components[0, 0, 0].global_scaling_const, + viscosity_mu_name=viscosity_mu_name) + + @override + def __str__(self) -> str: + return ( + f"StressletKnl{self.dim}D({self.viscosity_mu_name})") + + @override + def __reduce__(self): + return (type(self), (self.dim, self.viscosity_mu_name)) + + @property + @override + def shape(self) -> tuple[int, ...]: + return (self.dim, self.dim, self.dim) + + +@dataclass(frozen=True, repr=False) +class BrinkmanSystemKernel(ExpressionSystemKernel, ABC): + r"""A kernel for the Brinkman equations.""" + + mapper_method: ClassVar[str] = "map_brinkman_system_kernel" + + viscosity_mu_name: str + """The argument name to use for the dynamic viscosity when generating + functions to evaluate this kernel. + """ + darcy_impermeability_name: str + """The argument name to use for the Darcy impermeability when generating + functions to evaluate this kernel. + """ + + @property + @override + def is_complex_valued(self) -> bool: + # FIXME: 2D uses Hankel functions (complex-valued); 3D uses real exp + return self.dim == 2 + + @override + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: + from sumpy.expansion.diff_op import ( + concat, + divergence, + gradient, + laplacian, + make_identity_diff_op, + ) + + mu = sym.Symbol(self.viscosity_mu_name) + k = sym.Symbol(self.darcy_impermeability_name) + + u_and_p = make_identity_diff_op(self.dim, self.dim + 1) + u = u_and_p[:self.dim] + p = u_and_p[self.dim] + + return concat(mu * (laplacian(u) - k**2 * u) - gradient(p), divergence(u)) + + +@dataclass(frozen=True, repr=False) +class BrinkmanletSystemKernel(BrinkmanSystemKernel): + r"""A kernel for the Brinkman equations. + + This kernel uses :class:`BrinkmanletKernel` for its components. + """ + + mapper_method: ClassVar[str] = "map_brinkmanlet_system_kernel" + + def __init__(self, + dim: int, + viscosity_mu_name: str = "mu", + darcy_impermeability_name: str = "k") -> None: + components = {} + + for i in range(dim): + for j in range(i, dim): + components[i, j] = components[j, i] = BrinkmanletKernel( + dim, i, j, + viscosity_mu_name=viscosity_mu_name, + darcy_impermeability_name=darcy_impermeability_name, + ) + + super().__init__(dim=dim, + components=constantdict(components), + global_scaling_const=components[0, 0].global_scaling_const, + viscosity_mu_name=viscosity_mu_name, + darcy_impermeability_name=darcy_impermeability_name) + + @override + def __str__(self) -> str: + return ( + f"BrinkmanletKnl{self.dim}D" + f"({self.viscosity_mu_name}, {self.darcy_impermeability_name})") + + @override + def __reduce__(self): + return ( + type(self), + (self.dim, self.viscosity_mu_name, self.darcy_impermeability_name)) + + @property + @override + def shape(self) -> tuple[int, ...]: + return (self.dim, self.dim) + + +@dataclass(frozen=True, repr=False) +class BrinkmanStressSystemKernel(BrinkmanSystemKernel): + r"""A kernel for the Brinkman equations. + + This kernel uses :class:`BrinkmanStressKernel` for its components. + """ + + mapper_method: ClassVar[str] = "map_brinkman_stress_system_kernel" + + def __init__(self, + dim: int, + viscosity_mu_name: str = "mu", + darcy_impermeability_name: str = "k") -> None: + components = {} + + from itertools import product + for i, j, k in product(range(dim), repeat=3): + if j >= i and k >= j: + components[i, j, k] = BrinkmanStressKernel( + dim, i, j, k, + viscosity_mu_name=viscosity_mu_name, + darcy_impermeability_name=darcy_impermeability_name, + ) + else: + s = cast("tuple[int, int, int]", tuple(sorted([i, j, k]))) + components[i, j, k] = components[s] + + super().__init__(dim=dim, + components=constantdict(components), + global_scaling_const=components[0, 0, 0].global_scaling_const, + viscosity_mu_name=viscosity_mu_name, + darcy_impermeability_name=darcy_impermeability_name) + + @override + def __str__(self) -> str: + return ( + f"BrinkmanStressKnl{self.dim}D" + f"({self.viscosity_mu_name}, {self.darcy_impermeability_name})") + + @override + def __reduce__(self): + return ( + type(self), + (self.dim, self.viscosity_mu_name, self.darcy_impermeability_name)) + + @property + @override + def shape(self) -> tuple[int, ...]: + return (self.dim, self.dim, self.dim) # }}} @@ -1538,12 +2025,12 @@ def get_pde_as_diff_op(self): # {{{ a kernel defined as wrapping another one--e.g., derivatives @dataclass(frozen=True) -class KernelWrapper(Kernel, ABC): - inner_kernel: Kernel +class KernelWrapper(ScalarKernel, ABC): + inner_kernel: ScalarKernel """The kernel that is being wrapped (to take a derivative of, etc.).""" - def __init__(self, inner_kernel: Kernel) -> None: - Kernel.__init__(self, inner_kernel.dim) + def __init__(self, inner_kernel: ScalarKernel) -> None: + ScalarKernel.__init__(self, inner_kernel.dim) object.__setattr__(self, "inner_kernel", inner_kernel) @property @@ -1552,7 +2039,7 @@ def is_complex_valued(self) -> bool: return self.inner_kernel.is_complex_valued @override - def get_base_kernel(self) -> Kernel: + def get_base_kernel(self) -> ScalarKernel: return self.inner_kernel.get_base_kernel() @override @@ -1602,7 +2089,7 @@ def get_source_args(self) -> Sequence[KernelArgument]: return self.inner_kernel.get_source_args() @override - def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel: + def replace_base_kernel(self, new_base_kernel: ScalarKernel) -> ScalarKernel: raise NotImplementedError( f"'replace_base_kernel' is not implemented for '{type(self).__name__}'") @@ -1620,7 +2107,7 @@ def get_derivative_taker(self, # {{{ derivatives class DerivativeBase(KernelWrapper, ABC): - """Bases: :class:`Kernel` + """Bases: :class:`ScalarKernel` .. autoattribute:: inner_kernel .. automethod:: replace_inner_kernel @@ -1631,11 +2118,11 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: return self.inner_kernel.get_pde_as_diff_op() @abstractmethod - def replace_inner_kernel(self, new_inner_kernel: Kernel) -> Kernel: + def replace_inner_kernel(self, new_inner_kernel: ScalarKernel) -> ScalarKernel: """Replace the inner kernel of this wrapper. - This is essentially the same as :meth:`Kernel.replace_base_kernel`, but it does - not recurse. + This is essentially the same as :meth:`ScalarKernel.replace_base_kernel`, + but it does not recurse. """ @@ -1650,7 +2137,7 @@ class AxisSourceDerivative(DerivativeBase): axis: int """Direction axis for the source derivative.""" - def __init__(self, axis: int, inner_kernel: Kernel) -> None: + def __init__(self, axis: int, inner_kernel: ScalarKernel) -> None: super().__init__(inner_kernel) object.__setattr__(self, "axis", axis) @@ -1672,12 +2159,12 @@ def get_derivative_coeff_dict_at_source( return result @override - def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel: + def replace_base_kernel(self, new_base_kernel: ScalarKernel) -> ScalarKernel: return type(self)(self.axis, self.inner_kernel.replace_base_kernel(new_base_kernel)) @override - def replace_inner_kernel(self, new_inner_kernel: Kernel) -> Kernel: + def replace_inner_kernel(self, new_inner_kernel: ScalarKernel) -> ScalarKernel: return type(self)(self.axis, new_inner_kernel) @@ -1692,7 +2179,7 @@ class AxisTargetDerivative(DerivativeBase): axis: int - def __init__(self, axis: int, inner_kernel: Kernel) -> None: + def __init__(self, axis: int, inner_kernel: ScalarKernel) -> None: super().__init__(inner_kernel) object.__setattr__(self, "axis", axis) @@ -1730,12 +2217,12 @@ def postprocess_at_target( + inner_expr.diff(target_vec[self.axis])) @override - def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel: + def replace_base_kernel(self, new_base_kernel: ScalarKernel) -> ScalarKernel: return type(self)(self.axis, self.inner_kernel.replace_base_kernel(new_base_kernel)) @override - def replace_inner_kernel(self, new_inner_kernel: Kernel) -> Kernel: + def replace_inner_kernel(self, new_inner_kernel: ScalarKernel) -> ScalarKernel: return type(self)(self.axis, new_inner_kernel) @@ -1783,7 +2270,9 @@ class DirectionalDerivative(DerivativeBase): dir_vec_name: str """Name of the vector used for the direction.""" - def __init__(self, inner_kernel: Kernel, dir_vec_name: str | None = None) -> None: + def __init__(self, + inner_kernel: ScalarKernel, + dir_vec_name: str | None = None) -> None: if dir_vec_name is None: dir_vec_name = f"{self.directional_kind}_derivative_dir" @@ -1796,13 +2285,13 @@ def __str__(self) -> str: return fr"{self.dir_vec_name}·∇_{d} {self.inner_kernel}" @override - def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel: + def replace_base_kernel(self, new_base_kernel: ScalarKernel) -> ScalarKernel: return type(self)( self.inner_kernel.replace_base_kernel(new_base_kernel), dir_vec_name=self.dir_vec_name) @override - def replace_inner_kernel(self, new_inner_kernel: Kernel) -> Kernel: + def replace_inner_kernel(self, new_inner_kernel: ScalarKernel) -> ScalarKernel: return type(self)(new_inner_kernel, dir_vec_name=self.dir_vec_name) @@ -1826,8 +2315,7 @@ def transform(expr: Expression) -> Expression: def get_derivative_coeff_dict_at_source( self, expr_dict: DerivativeCoeffDict, ) -> DerivativeCoeffDict: - from sumpy.symbolic import make_sym_vector as make_sympy_vector - dir_vec = make_sympy_vector(self.dir_vec_name, self.dim) + dir_vec = sym.make_sym_vector(self.dir_vec_name, self.dim) expr_dict = self.inner_kernel.get_derivative_coeff_dict_at_source( expr_dict) @@ -1862,7 +2350,7 @@ def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationU @dataclass(frozen=True) class TargetPointMultiplier(KernelWrapper): - """Bases: :class:`Kernel` + """Bases: :class:`ScalarKernel` Wraps a kernel :math:`G(x, y)` and outputs :math:`x_j G(x, y)` where :math:`x, y` are targets and sources respectively. @@ -1876,7 +2364,7 @@ class TargetPointMultiplier(KernelWrapper): axis: int """Coordinate axis with which to multiply the kernel.""" - def __init__(self, axis: int, inner_kernel: Kernel) -> None: + def __init__(self, axis: int, inner_kernel: ScalarKernel) -> None: super().__init__(inner_kernel) object.__setattr__(self, "axis", axis) @@ -1933,11 +2421,11 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: raise NotImplementedError("no PDE is known") @override - def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel: + def replace_base_kernel(self, new_base_kernel: ScalarKernel) -> ScalarKernel: return type(self)(self.axis, self.inner_kernel.replace_base_kernel(new_base_kernel)) - def replace_inner_kernel(self, new_inner_kernel: Kernel) -> Kernel: + def replace_inner_kernel(self, new_inner_kernel: ScalarKernel) -> ScalarKernel: return type(self)(self.axis, new_inner_kernel) # }}} @@ -1952,17 +2440,17 @@ class KernelMapper(Generic[ResultT]): """ .. automethod:: __call__ """ - def rec(self, kernel: Kernel) -> ResultT: + def rec(self, kernel: ScalarKernel) -> ResultT: try: method = cast( - "Callable[[Kernel], ResultT]", + "Callable[[ScalarKernel], ResultT]", getattr(self, kernel.mapper_method)) except AttributeError as err: raise RuntimeError(f"{type(self)} cannot handle {type(kernel)}") from err else: return method(kernel) - def __call__(self, kernel: Kernel) -> ResultT: + def __call__(self, kernel: ScalarKernel) -> ResultT: return self.rec(kernel) @@ -1992,44 +2480,33 @@ def map_target_point_multiplier( return self.rec(kernel.inner_kernel) -class KernelIdentityMapper(KernelMapper[Kernel]): - def map_expression_kernel(self, kernel: ExpressionKernel) -> Kernel: +class KernelIdentityMapper(KernelMapper[ScalarKernel]): + def map_expression_kernel(self, kernel: ExpressionKernel) -> ScalarKernel: return kernel - map_laplace_kernel: \ - Callable[[Self, LaplaceKernel], Kernel] = map_expression_kernel - map_biharmonic_kernel: \ - Callable[[Self, BiharmonicKernel], Kernel] = map_expression_kernel - map_helmholtz_kernel: \ - Callable[[Self, HelmholtzKernel], Kernel] = map_expression_kernel - map_yukawa_kernel: \ - Callable[[Self, YukawaKernel], Kernel] = map_expression_kernel - map_elasticity_kernel: \ - Callable[[Self, ElasticityKernel], Kernel] = map_expression_kernel - map_line_of_compression_kernel: \ - Callable[[Self, LineOfCompressionKernel], Kernel] = map_expression_kernel - map_stokeslet_kernel: \ - Callable[[Self, StokesletKernel], Kernel] = map_expression_kernel - map_stresslet_kernel: \ - Callable[[Self, StressletKernel], Kernel] = map_expression_kernel - map_brinkmanlet_kernel: \ - Callable[[Self, BrinkmanletKernel], Kernel] = map_expression_kernel - map_brinkman_stress_kernel: \ - Callable[[Self, BrinkmanStressKernel], Kernel] = map_expression_kernel - map_heat_kernel: \ - Callable[[Self, HeatKernel], Kernel] = map_expression_kernel - - def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> Kernel: + map_laplace_kernel: Callable[[Self, LaplaceKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_biharmonic_kernel: Callable[[Self, BiharmonicKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_helmholtz_kernel: Callable[[Self, HelmholtzKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_yukawa_kernel: Callable[[Self, YukawaKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_elasticity_kernel: Callable[[Self, ElasticityKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_line_of_compression_kernel: Callable[[Self, LineOfCompressionKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_stokeslet_kernel: Callable[[Self, StokesletKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_stresslet_kernel: Callable[[Self, StressletKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_brinkmanlet_kernel: Callable[[Self, BrinkmanletKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_brinkman_stress_kernel: Callable[[Self, BrinkmanStressKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_heat_kernel: Callable[[Self, HeatKernel], ScalarKernel] = map_expression_kernel + + def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> ScalarKernel: return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel)) - def map_axis_source_derivative(self, kernel: AxisSourceDerivative) -> Kernel: + def map_axis_source_derivative(self, kernel: AxisSourceDerivative) -> ScalarKernel: return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel)) - def map_target_point_multiplier(self, kernel: TargetPointMultiplier) -> Kernel: + def map_target_point_multiplier(self, kernel: TargetPointMultiplier) -> ScalarKernel: # noqa: E501 return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel)) def map_directional_source_derivative( - self, kernel: DirectionalSourceDerivative) -> Kernel: + self, kernel: DirectionalSourceDerivative) -> ScalarKernel: return type(kernel)(self.rec(kernel.inner_kernel), dir_vec_name=kernel.dir_vec_name) @@ -2038,7 +2515,7 @@ class AxisSourceDerivativeRemover(KernelIdentityMapper): """Removes all axis source derivatives from the kernel.""" @override - def map_axis_source_derivative(self, kernel: AxisSourceDerivative) -> Kernel: + def map_axis_source_derivative(self, kernel: AxisSourceDerivative) -> ScalarKernel: return self.rec(kernel.inner_kernel) @@ -2046,7 +2523,7 @@ class AxisTargetDerivativeRemover(KernelIdentityMapper): """Removes all axis target derivatives from the kernel.""" @override - def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> Kernel: + def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> ScalarKernel: return self.rec(kernel.inner_kernel) @@ -2059,7 +2536,7 @@ class SourceDerivativeRemover(AxisSourceDerivativeRemover): @override def map_directional_source_derivative( - self, kernel: DirectionalSourceDerivative) -> Kernel: + self, kernel: DirectionalSourceDerivative) -> ScalarKernel: return self.rec(kernel.inner_kernel) @@ -2067,7 +2544,8 @@ class TargetTransformationRemover(TargetDerivativeRemover): """Removes all target transformations from the kernel.""" @override - def map_target_point_multiplier(self, kernel: TargetPointMultiplier) -> Kernel: + def map_target_point_multiplier( + self, kernel: TargetPointMultiplier) -> ScalarKernel: return self.rec(kernel.inner_kernel) diff --git a/sumpy/p2p.py b/sumpy/p2p.py index ad80cdd7..2a320766 100644 --- a/sumpy/p2p.py +++ b/sumpy/p2p.py @@ -72,9 +72,9 @@ class P2PBase(KernelCacheMixin, KernelComputation): def __init__(self, target_kernels, exclude_self, strength_usage=None, value_dtypes=None, name=None, source_kernels=None): """ - :arg target_kernels: list of :class:`sumpy.kernel.Kernel` instances + :arg target_kernels: list of :class:`~sumpy.kernel.ScalarKernel` instances with only target derivatives. - :arg source_kernels: list of :class:`sumpy.kernel.Kernel` instances + :arg source_kernels: list of :class:`~sumpy.kernel.ScalarKernel` instances with only source derivatives. :arg strength_usage: A list of integers indicating which expression uses which source strength indicator. This implicitly specifies the diff --git a/sumpy/test/test_fmm.py b/sumpy/test/test_fmm.py index 45455f72..4936f09b 100644 --- a/sumpy/test/test_fmm.py +++ b/sumpy/test/test_fmm.py @@ -57,8 +57,8 @@ from sumpy.kernel import ( BiharmonicKernel, HelmholtzKernel, - Kernel, LaplaceKernel, + ScalarKernel, YukawaKernel, ) @@ -105,7 +105,7 @@ ]) def test_sumpy_fmm( actx_factory: ArrayContextFactory, - knl: Kernel, + knl: ScalarKernel, local_expn_class, mpole_expn_class, order_varies_with_level, @@ -140,7 +140,7 @@ def test_sumpy_fmm( def _test_sumpy_fmm( actx_factory: ArrayContextFactory, - knl: Kernel, + knl: ScalarKernel, local_expn_class, mpole_expn_class, order_varies_with_level, diff --git a/sumpy/test/test_kernels.py b/sumpy/test/test_kernels.py index 44f13c35..f1334ea6 100644 --- a/sumpy/test/test_kernels.py +++ b/sumpy/test/test_kernels.py @@ -67,10 +67,10 @@ DirectionalSourceDerivative, ElasticityKernel, HelmholtzKernel, - Kernel, LaplaceKernel, LineOfCompressionKernel, OneKernel, + ScalarKernel, StokesletKernel, StressletKernel, YukawaKernel, @@ -152,7 +152,7 @@ def test_p2p(actx_factory: ArrayContextFactory, exclude_self): ]) def test_p2e_multiple( actx_factory: ArrayContextFactory, - base_knl: Kernel, + base_knl: ScalarKernel, expn_class): order = 4 actx = actx_factory() @@ -844,7 +844,7 @@ def test_m2m_compressed_error_helmholtz(actx_factory: ArrayContextFactory, dim, @pytest.mark.parametrize("dim", [2, 3]) def test_jump( actx_factory: ArrayContextFactory, - kernel_cls: type[Kernel], + kernel_cls: type[ScalarKernel], kernel_kwargs: dict[str, float], dim: int, order: int = 3, @@ -907,7 +907,7 @@ def test_jump( # {{{ test_pickle @memoize_on_first_arg -def get_kernel_name_for_test(knl: Kernel) -> Callable[[str], str]: +def get_kernel_name_for_test(knl: ScalarKernel) -> Callable[[str], str]: return lambda prefix: f"{prefix}: {type(knl).__name__}" @@ -930,7 +930,7 @@ def get_kernel_name_for_test(knl: Kernel) -> Callable[[str], str]: StressletKernel(2, 0, 0, 0), YukawaKernel(2, yukawa_lambda_name="lambda"), ]) -def test_pickle(knl: Kernel) -> None: +def test_pickle(knl: ScalarKernel) -> None: import pickle result = pickle.dumps(knl) diff --git a/sumpy/test/test_l2l_coeffs.py b/sumpy/test/test_l2l_coeffs.py index 482ae95f..cfc08867 100644 --- a/sumpy/test/test_l2l_coeffs.py +++ b/sumpy/test/test_l2l_coeffs.py @@ -47,8 +47,8 @@ from sumpy.kernel import ( BiharmonicKernel, HelmholtzKernel, - Kernel, LaplaceKernel, + ScalarKernel, YukawaKernel, ) from sumpy.tools import add_mi, build_matrix, mi_factorial, mi_power @@ -73,7 +73,7 @@ ]) def test_l2l_coefficient_differences( actx_factory: ArrayContextFactory, - knl: Kernel, + knl: ScalarKernel, extra_kwargs: Mapping[str, float], verbose: bool = True, ): diff --git a/sumpy/test/test_m2m_coeffs.py b/sumpy/test/test_m2m_coeffs.py index a44f1030..07acdd63 100644 --- a/sumpy/test/test_m2m_coeffs.py +++ b/sumpy/test/test_m2m_coeffs.py @@ -50,8 +50,8 @@ from sumpy.kernel import ( BiharmonicKernel, HelmholtzKernel, - Kernel, LaplaceKernel, + ScalarKernel, YukawaKernel, ) from sumpy.tools import add_mi, build_matrix, mi_factorial, mi_power @@ -76,7 +76,7 @@ ]) def test_m2m_coefficient_differences( actx_factory: ArrayContextFactory, - knl: Kernel, + knl: ScalarKernel, extra_kwargs: Mapping[str, float], verbose: bool = True, ): diff --git a/sumpy/test/test_misc.py b/sumpy/test/test_misc.py index 04f05169..d5eff0ed 100644 --- a/sumpy/test/test_misc.py +++ b/sumpy/test/test_misc.py @@ -59,16 +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, ) @@ -87,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() @@ -226,7 +232,7 @@ class FakeTree: @pytest.mark.parametrize("knl", [ LaplaceKernel(2), HelmholtzKernel(2), LaplaceKernel(3), HelmholtzKernel(3)]) -def test_order_finder(knl: Kernel) -> None: +def test_order_finder(knl: ScalarKernel) -> None: from sumpy.expansion.level_to_order import SimpleExpansionOrderFinder ofind = SimpleExpansionOrderFinder(1e-5) @@ -790,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)' diff --git a/sumpy/test/test_target_deriv.py b/sumpy/test/test_target_deriv.py index dc03c431..f5d0648e 100644 --- a/sumpy/test/test_target_deriv.py +++ b/sumpy/test/test_target_deriv.py @@ -43,7 +43,7 @@ _acf, # pyright: ignore[reportUnusedImport] ) from sumpy.expansion.local import LineTaylorLocalExpansion -from sumpy.kernel import AxisTargetDerivative, Kernel, LaplaceKernel +from sumpy.kernel import AxisTargetDerivative, LaplaceKernel, ScalarKernel from sumpy.test.geometries import make_starfish @@ -58,7 +58,7 @@ @pytest.mark.parametrize("knl", [LaplaceKernel(2)]) def test_lpot_dx_jump_relation_convergence( actx_factory: ArrayContextFactory, - knl: Kernel): + knl: ScalarKernel): """Test convergence of jump relations for single layer potential derivatives.""" actx = actx_factory() diff --git a/sumpy/tools.py b/sumpy/tools.py index 755f13ad..5361a183 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -60,7 +60,7 @@ from sumpy.assignment_collection import SymbolicAssignmentCollection from sumpy.expansion import ExpansionBase - from sumpy.kernel import Kernel, KernelArgument + from sumpy.kernel import KernelArgument, ScalarKernel logger = logging.getLogger(__name__) @@ -134,7 +134,7 @@ """ -KernelLike: TypeAlias = "Kernel | ExpansionBase" +KernelLike: TypeAlias = "ScalarKernel | ExpansionBase" # {{{ multi_index helpers @@ -285,16 +285,16 @@ class KernelComputation(ABC): """ def __init__(self, - target_kernels: list[Kernel], - source_kernels: list[Kernel], + target_kernels: list[ScalarKernel], + source_kernels: list[ScalarKernel], strength_usage: list[int] | None = None, value_dtypes: list[numpy.dtype[Any]] | None = None, name: str | None = None) -> None: """ - :arg target_kernels: list of :class:`~sumpy.kernel.Kernel` instances, + :arg target_kernels: list of :class:`~sumpy.kernel.ScalarKernel` instances, with :class:`sumpy.kernel.AxisTargetDerivative` as the outermost kernel wrappers, if present. - :arg source_kernels: list of :class:`~sumpy.kernel.Kernel` instances + :arg source_kernels: list of :class:`~sumpy.kernel.ScalarKernel` instances with :class:`~sumpy.kernel.DirectionalSourceDerivative` as the outermost kernel wrappers, if present. :arg strength_usage: list of integers indicating which expression diff --git a/sumpy/toys.py b/sumpy/toys.py index 6d77405f..d41bb47a 100644 --- a/sumpy/toys.py +++ b/sumpy/toys.py @@ -50,7 +50,7 @@ MultipoleExpansionFactory, ) from sumpy.expansion.m2l import M2LTranslationClassFactoryBase - from sumpy.kernel import Kernel + from sumpy.kernel import ScalarKernel from sumpy.visualization import FieldPlotter @@ -107,8 +107,8 @@ class ToyContext: .. automethod:: __init__ """ - kernel: Kernel - no_target_deriv_kernel: Kernel + kernel: ScalarKernel + no_target_deriv_kernel: ScalarKernel mpole_expn_class: MultipoleExpansionFactory local_expn_class: LocalExpansionFactory @@ -118,7 +118,7 @@ class ToyContext: extra_source_and_kernel_kwargs: Mapping[str, object] def __init__(self, - kernel: Kernel, + kernel: ScalarKernel, mpole_expn_class: MultipoleExpansionFactory | None = None, local_expn_class: LocalExpansionFactory | None = None, expansion_factory: ExpansionFactoryBase | None = None,