diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index b8384a04..a26398c7 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -9103,6 +9103,14 @@ "lineCount": 1 } }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 15, + "endColumn": 24, + "lineCount": 1 + } + }, { "code": "reportAny", "range": { @@ -9135,6 +9143,22 @@ "lineCount": 1 } }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 21, + "lineCount": 1 + } + }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 17, + "lineCount": 1 + } + }, { "code": "reportUnknownMemberType", "range": { @@ -9223,6 +9247,94 @@ "lineCount": 1 } }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 51, + "lineCount": 1 + } + }, + { + "code": "reportIncompatibleMethodOverride", + "range": { + "startColumn": 8, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportReturnType", + "range": { + "startColumn": 15, + "endColumn": 78, + "lineCount": 1 + } + }, + { + "code": "reportIncompatibleMethodOverride", + "range": { + "startColumn": 8, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 22, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportIncompatibleMethodOverride", + "range": { + "startColumn": 8, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 22, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportIncompatibleMethodOverride", + "range": { + "startColumn": 8, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 22, + "endColumn": 66, + "lineCount": 1 + } + }, + { + "code": "reportIncompatibleMethodOverride", + "range": { + "startColumn": 8, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 22, + "endColumn": 66, + "lineCount": 1 + } + }, { "code": "reportIncompatibleMethodOverride", "range": { @@ -9422,6 +9534,14 @@ "endColumn": 59, "lineCount": 1 } + }, + { + "code": "reportAny", + "range": { + "startColumn": 4, + "endColumn": 15, + "lineCount": 1 + } } ], "./sumpy/p2e.py": [ @@ -21092,26 +21212,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 +21284,10 @@ } }, { - "code": "reportUnknownMemberType", + "code": "reportAny", "range": { - "startColumn": 50, - "endColumn": 67, + "startColumn": 45, + "endColumn": 46, "lineCount": 1 } }, @@ -22610,6 +22714,86 @@ "endColumn": 56, "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 + } } ], "./sumpy/test/test_qbx.py": [ @@ -27295,22 +27479,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 8, - "endColumn": 15, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 8, - "endColumn": 15, - "lineCount": 1 - } - }, { "code": "reportAny", "range": { diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a22d16eb..5d2afb94 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,7 +42,7 @@ jobs: curl -L -O https://tiker.net/ci-support-v0 . ./ci-support-v0 build_py_project_in_conda_env - cipip install pytest pyfmmlib scipy scipy-stubs matplotlib pyvisfile optype + cipip install pytest pyfmmlib scipy scipy-stubs matplotlib pyvisfile 'optype<0.18' cipip install basedpyright basedpyright 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..c7e14f26 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -28,6 +28,7 @@ from dataclasses import dataclass from typing import ( TYPE_CHECKING, + Any, ClassVar, Generic, Literal, @@ -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,7 +54,6 @@ ExprDerivativeTaker, diff_derivative_coeff_dict, ) -from sumpy.symbolic import SpatialConstant, pymbolic_real_norm_2 if TYPE_CHECKING: @@ -69,7 +69,9 @@ ---------------- .. autoclass:: KernelArgument -.. autoclass:: Kernel +.. autoclass:: ScalarKernel + :show-inheritance: +.. autoclass:: SystemKernel :show-inheritance: Symbolic kernels @@ -79,43 +81,74 @@ :show-inheritance: :members: mapper_method -PDE kernels ------------ +.. 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 + +.. autoclass:: StokesletComponentKernel :show-inheritance: :members: mapper_method -.. autoclass:: StressletKernel +.. autoclass:: StressletComponentKernel :show-inheritance: :members: mapper_method -.. autoclass:: ElasticityKernel + +.. autoclass:: ElasticityComponentKernel :show-inheritance: :members: mapper_method .. autoclass:: LineOfCompressionKernel :show-inheritance: :members: mapper_method -.. autoclass:: BrinkmanletKernel + +.. autoclass:: BrinkmanletComponentKernel :show-inheritance: :members: mapper_method -.. autoclass:: BrinkmanStressKernel +.. autoclass:: BrinkmanStressComponentKernel :show-inheritance: :members: mapper_method .. autoclass:: HeatKernel :show-inheritance: :members: mapper_method +System PDE Kernels +------------------ + +.. autoclass:: ElasticitySystemKernel + :show-inheritance: + :members: mapper_method + +.. autoclass:: StokesletSystemKernel + :show-inheritance: + :members: mapper_method +.. autoclass:: StressletSystemKernel + :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,9 +231,23 @@ def name(self) -> str: # {{{ basic kernel interface -@dataclass(frozen=True) -class Kernel(ABC): - """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 + + +@dataclass(frozen=True, repr=False) +class ScalarKernel(ABC): + """Scalar kernel interface. .. autoattribute:: mapper_method .. autoattribute:: is_translation_invariant @@ -215,6 +262,8 @@ class Kernel(ABC): .. automethod:: get_expression .. automethod:: postprocess_at_source .. automethod:: postprocess_at_target + .. automethod:: get_pde_system_kernel + .. automethod:: get_global_scaling_const .. automethod:: get_args .. automethod:: get_source_args @@ -249,19 +298,27 @@ def __repr__(self) -> str: return f"{type(self).__name__}({', '.join(args)})" - def get_base_kernel(self) -> Kernel: + 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_pde_system_kernel(self) -> tuple[SystemKernel, tuple[int, ...]]: + """ + :returns: if the kernel is a component kernel of a :class:`SystemKernel`, + this returns the system kernel and the index of the kernel in that + system. + """ + raise TypeError(f"kernel {type(self)} is not part of a system") + def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit: """Apply some changes (such as registering function manglers) to the kernel. @@ -343,7 +400,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 @@ -411,11 +468,83 @@ def get_source_args(self) -> Sequence[KernelArgument]: """ return [] + +@dataclass(frozen=True, repr=False) +class SystemKernel(ABC): + """A kernel representing a vector PDE. + + .. autoattribute:: dim + .. autoproperty:: ndim + .. autoproperty:: shape + + .. automethod:: __getitem__ + .. automethod:: get_expression + .. automethod:: get_pde_as_diff_op + """ + + dim: int + """Dimension of the space the kernel is defined in.""" + + @property + def ndim(self) -> int: + """The number of indices in the kernel tensor.""" + return len(self.shape) + + @property + @abstractmethod + def shape(self) -> tuple[int, ...]: + """The shape of the kernel tensor.""" + + def __getitem__(self, index: tuple[int, ...], /) -> ScalarKernel: + """ + :returns: the scalar kernel at *index*. + """ + if len(index) != self.ndim: + raise IndexError( + f"incorrect index size for kernel: kernel is {self.ndim}-dimensional: " + f"{index} given" + ) + + if any(not 0 <= i < n for i, n in zip(index, self.shape, strict=True)): + raise IndexError( + f"index {index} is out of bounds for kernel with shape {self.shape}" + ) + + return self.get_scalar_component(*index) + + @abstractmethod + def get_scalar_component(self, *args: int) -> ScalarKernel: + """ + :returns: the scalar kernel the indices given by *args*. + """ + + def get_expression(self, dist_vec: sym.Matrix) -> obj_array.ObjectArrayND[sym.Expr]: + """ + :returns: a :mod:`sympy` expression for each component the kernel. + """ + from pytools import ndindex + + result = np.empty(self.shape, dtype=object) + for i in ndindex(result.shape): + result[i] = self[i].get_expression(dist_vec) + + return result + + @abstractmethod + def get_pde_as_diff_op(self) -> LinearPDESystemOperator: + """ + :returns: the PDE that this kernel satisfies + (see :meth:`ScalarKernel.get_pde_as_diff_op` as the scalar alternative). + """ + 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 +575,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( @@ -500,6 +626,8 @@ def is_complex_valued(self) -> bool: one_kernel_2d = OneKernel(2) one_kernel_3d = OneKernel(3) +# }}} + # {{{ PDE kernels @@ -516,11 +644,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 +697,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 +771,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 +817,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 +859,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 +870,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 +913,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( @@ -810,7 +936,7 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: @dataclass(frozen=True, repr=False) -class ElasticityKernel(ExpressionKernel): +class ElasticityComponentKernel(ExpressionKernel): r"""A kernel for the linear elasticity (Navier-Cauchy) equations (see e.g. Section 2.2 in [Hsiao2008]_). @@ -858,11 +984,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: @@ -910,6 +1036,14 @@ def get_args(self) -> Sequence[KernelArgument]: KernelArgument(loopy_arg=lp.ValueArg(self.poisson_ratio_name, np.float64)), ] + @override + def get_pde_system_kernel(self) -> tuple[SystemKernel, tuple[int, ...]]: + return ElasticitySystemKernel( + self.dim, + viscosity_mu_name=self.viscosity_mu_name, + poisson_ratio_name=self.poisson_ratio_name + ), (self.icomp, self.jcomp) + @override def get_pde_as_diff_op(self) -> LinearPDESystemOperator: from sumpy.expansion.diff_op import laplacian, make_identity_diff_op @@ -919,7 +1053,7 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: @dataclass(frozen=True, repr=False) -class StokesletKernel(ExpressionKernel): +class StokesletComponentKernel(ExpressionKernel): r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_). .. math:: @@ -961,10 +1095,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: @@ -1007,6 +1141,13 @@ def get_args(self) -> Sequence[KernelArgument]: ) ] + @override + def get_pde_system_kernel(self) -> tuple[SystemKernel, tuple[int, ...]]: + return StokesletSystemKernel( + self.dim, + viscosity_mu_name=self.viscosity_mu_name, + ), (self.icomp, self.jcomp) + @override def get_pde_as_diff_op(self) -> LinearPDESystemOperator: from sumpy.expansion.diff_op import laplacian, make_identity_diff_op @@ -1016,7 +1157,7 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: @dataclass(frozen=True, repr=False) -class StressletKernel(ExpressionKernel): +class StressletComponentKernel(ExpressionKernel): r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_). .. math:: @@ -1025,8 +1166,9 @@ class StressletKernel(ExpressionKernel): -P_j \delta_{ik} + \mu (\partial_k K_{ij} + \partial_i K_{kj}) - where the two-index :math:`K_{ij}` is the :class:`~sumpy.kernel.StokesletKernel`. - This kernel is often called the Stresslet and it represents the stress tensor. + where the two-index :math:`K_{ij}` is the + :class:`~sumpy.kernel.StokesletComponentKernel`. This kernel is often + called the Stresslet and it represents the stress tensor. .. autoattribute:: icomp .. autoattribute:: jcomp @@ -1060,7 +1202,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 @@ -1093,7 +1235,7 @@ def is_complex_valued(self) -> bool: @memoize_method @override def get_args(self) -> Sequence[KernelArgument]: - # NOTE: this is not used, but kept for symmetry with the StokesletKernel + # NOTE: this is not used, kept for symmetry with the StokesletComponentKernel return [ KernelArgument(loopy_arg=lp.ValueArg(self.viscosity_mu_name, np.float64)) ] @@ -1104,6 +1246,13 @@ def __str__(self) -> str: f"StressletKnl{self.dim}D_{self.icomp}{self.jcomp}{self.kcomp}" f"({self.viscosity_mu_name})") + @override + def get_pde_system_kernel(self) -> tuple[SystemKernel, tuple[int, ...]]: + return StressletSystemKernel( + self.dim, + viscosity_mu_name=self.viscosity_mu_name, + ), (self.icomp, self.jcomp, self.kcomp) + @override def get_pde_as_diff_op(self) -> LinearPDESystemOperator: from sumpy.expansion.diff_op import laplacian, make_identity_diff_op @@ -1160,12 +1309,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 @@ -1214,7 +1363,7 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: @dataclass(frozen=True, repr=False) -class BrinkmanletKernel(ExpressionKernel): +class BrinkmanletComponentKernel(ExpressionKernel): r"""A kernel for the Brinkman equations. .. math:: @@ -1252,11 +1401,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 @@ -1327,17 +1476,26 @@ def get_args(self) -> Sequence[KernelArgument]: KernelArgument(loopy_arg=lp.ValueArg(self.darcy_impermeability_name, np.float64)), # noqa: E501 ] + @override + def get_pde_system_kernel(self) -> tuple[SystemKernel, tuple[int, ...]]: + return BrinkmanletSystemKernel( + self.dim, + viscosity_mu_name=self.viscosity_mu_name, + darcy_impermeability_name=self.darcy_impermeability_name, + ), (self.icomp, self.jcomp) + @override def get_pde_as_diff_op(self) -> LinearPDESystemOperator: from sumpy.expansion.diff_op import laplacian, make_identity_diff_op w = make_identity_diff_op(self.dim) k = sym.Symbol(self.darcy_impermeability_name) + return laplacian(laplacian(w) - k**2 * w) @dataclass(frozen=True, repr=False) -class BrinkmanStressKernel(ExpressionKernel): +class BrinkmanStressComponentKernel(ExpressionKernel): r"""A kernel for the Brinkman equations. .. math:: @@ -1346,8 +1504,9 @@ class BrinkmanStressKernel(ExpressionKernel): -p_j \delta_{ik} + \mu (\partial_k K_{ij} + \partial_i K_{jk}), - where the two-index :math:`K_{ij}` is the :class:`~sumpy.kernel.BrinkmanletKernel` - and :math:`P_j` is the pressure kernel. + where the two-index :math:`K_{ij}` is the + :class:`~sumpy.kernel.BrinkmanletComponentKernel` and :math:`P_j` is the + pressure kernel. """ mapper_method: ClassVar[str] = "map_brinkman_stress_kernel" @@ -1377,10 +1536,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 @@ -1464,6 +1623,14 @@ def get_args(self) -> Sequence[KernelArgument]: KernelArgument(loopy_arg=lp.ValueArg(self.darcy_impermeability_name, np.float64)), # noqa: E501 ] + @override + def get_pde_system_kernel(self) -> tuple[SystemKernel, tuple[int, ...]]: + return BrinkmanStressSystemKernel( + self.dim, + viscosity_mu_name=self.viscosity_mu_name, + darcy_impermeability_name=self.darcy_impermeability_name, + ), (self.icomp, self.jcomp, self.kcomp) + @override def get_pde_as_diff_op(self) -> LinearPDESystemOperator: from sumpy.expansion.diff_op import laplacian, make_identity_diff_op @@ -1475,61 +1642,387 @@ 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(SystemKernel): + r"""A kernel for the linear elasticity (Navier-Cauchy) equations + (see e.g. Section 2.2 in [Hsiao2008]_). + + This kernel uses :class:`ElasticityComponentKernel` for its components. + + .. autoattribute:: viscosity_mu_name + .. autoattribute:: poisson_ratio_name + """ + + viscosity_mu_name: str = "mu" + r"""The argument name to use for the dynamic viscosity :math:`\mu` when + generating functions to evaluate this kernel. + """ + poisson_ratio_name: str = "nu" + r"""The argument name to use for Poisson's ratio :math:`\nu` when generating + functions to evaluate this kernel. + """ + + @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 shape(self) -> tuple[int, ...]: + return (self.dim, self.dim) + + @memoize_method + @override + def get_scalar_component(self, i: int, j: int, /) -> ScalarKernel: + # NOTE: the kernel is (i, j) -> (j, i) symmetric + i, j = sorted([i, j]) + + return ElasticityComponentKernel( + self.dim, i, j, + viscosity_mu_name=self.viscosity_mu_name, + poisson_ratio_name=self.poisson_ratio_name, + ) + + @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 StokesletSystemKernel(SystemKernel): + r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_). + + This kernel uses :class:`StokesletComponentKernel` for its components. + + .. autoattribute:: viscosity_mu_name + """ + + viscosity_mu_name: str = "mu" + r"""The argument name to use for the dynamic viscosity :math:`\mu` when + generating functions to evaluate this kernel. + """ + + @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) + + @memoize_method + @override + def get_scalar_component(self, i: int, j: int, /) -> ScalarKernel: + # NOTE: the kernel is (i, j) -> (j, i) symmetric + i, j = sorted([i, j]) + + return StokesletComponentKernel( + self.dim, i, j, + viscosity_mu_name=self.viscosity_mu_name, + ) + + @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 StressletSystemKernel(SystemKernel): + r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_). + + This kernel uses :class:`StressletComponentKernel` for its components. + + .. autoattribute:: viscosity_mu_name + """ + + viscosity_mu_name: str = "mu" + r"""The argument name to use for the dynamic viscosity :math:`\mu` when + generating functions to evaluate this kernel. + """ + + @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) + + @memoize_method + @override + def get_scalar_component(self, i: int, j: int, k: int, /) -> ScalarKernel: + # NOTE: the kernel is (i, j, k) -> (j, i, k) symmetric + i, j, k = sorted([i, j, k]) + + return StressletComponentKernel( + self.dim, i, j, k, + viscosity_mu_name=self.viscosity_mu_name, + ) + + @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 BrinkmanletSystemKernel(SystemKernel): + r"""A kernel for the Brinkman equations. + + This kernel uses :class:`BrinkmanletComponentKernel` for its components. + + .. autoattribute:: viscosity_mu_name + .. autoattribute:: darcy_impermeability_name + """ + + viscosity_mu_name: str = "mu" + """The argument name to use for the dynamic viscosity when generating + functions to evaluate this kernel. + """ + darcy_impermeability_name: str = "k" + """The argument name to use for the Darcy impermeability when generating + functions to evaluate this kernel. + """ + + @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) + + @memoize_method + @override + def get_scalar_component(self, i: int, j: int, /) -> ScalarKernel: + # NOTE: the kernel is (i, j) -> (j, i) symmetric + i, j = sorted([i, j]) + + return BrinkmanletComponentKernel( + self.dim, i, j, + viscosity_mu_name=self.viscosity_mu_name, + darcy_impermeability_name=self.darcy_impermeability_name, + ) + + @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 BrinkmanStressSystemKernel(SystemKernel): + r"""A kernel for the Brinkman equations. + + This kernel uses :class:`BrinkmanStressComponentKernel` for its components. + + .. autoattribute:: viscosity_mu_name + .. autoattribute:: darcy_impermeability_name + """ + + viscosity_mu_name: str = "mu" + """The argument name to use for the dynamic viscosity when generating + functions to evaluate this kernel. + """ + darcy_impermeability_name: str = "k" + """The argument name to use for the Darcy impermeability when generating + functions to evaluate this kernel. + """ + + @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) + + @memoize_method + @override + def get_scalar_component(self, i: int, j: int, k: int, /) -> ScalarKernel: + # NOTE: the kernel is (i, j) -> (j, i, k) symmetric + i, j, k = sorted([i, j, k]) + + return BrinkmanStressComponentKernel( + self.dim, i, j, k, + viscosity_mu_name=self.viscosity_mu_name, + darcy_impermeability_name=self.darcy_impermeability_name, + ) + + @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)) # }}} @@ -1538,12 +2031,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 +2045,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 +2095,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 +2113,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 +2124,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 +2143,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 +2165,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 +2185,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 +2223,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 +2276,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 +2291,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 +2321,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 +2356,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 +2370,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 +2427,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 +2446,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 +2486,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, ElasticityComponentKernel], 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, StokesletComponentKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_stresslet_kernel: Callable[[Self, StressletComponentKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_brinkmanlet_kernel: Callable[[Self, BrinkmanletComponentKernel], ScalarKernel] = map_expression_kernel # noqa: E501 + map_brinkman_stress_kernel: Callable[[Self, BrinkmanStressComponentKernel], 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 +2521,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 +2529,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 +2542,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 +2550,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) @@ -2093,17 +2577,17 @@ def map_expression_kernel(self, kernel: ExpressionKernel) -> int: map_yukawa_kernel: \ Callable[[Self, YukawaKernel], int] = map_expression_kernel map_elasticity_kernel: \ - Callable[[Self, ElasticityKernel], int] = map_expression_kernel + Callable[[Self, ElasticityComponentKernel], int] = map_expression_kernel map_line_of_compression_kernel: \ Callable[[Self, LineOfCompressionKernel], int] = map_expression_kernel map_stokeslet_kernel: \ - Callable[[Self, StokesletKernel], int] = map_expression_kernel + Callable[[Self, StokesletComponentKernel], int] = map_expression_kernel map_stresslet_kernel: \ - Callable[[Self, StressletKernel], int] = map_expression_kernel + Callable[[Self, StressletComponentKernel], int] = map_expression_kernel map_brinkmanlet_kernel: \ - Callable[[Self, BrinkmanletKernel], int] = map_expression_kernel + Callable[[Self, BrinkmanletComponentKernel], int] = map_expression_kernel map_brinkman_stress_kernel: \ - Callable[[Self, BrinkmanStressKernel], int] = map_expression_kernel + Callable[[Self, BrinkmanStressComponentKernel], int] = map_expression_kernel map_heat_kernel: \ Callable[[Self, HeatKernel], int] = map_expression_kernel @@ -2118,4 +2602,35 @@ def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> int: # }}} +# {{{ deprecations + +# TODO: once these deprecations expire, rename +# ElasticitySystemKernel -> ElasticityKernel +# ... +_DEPRECATED_CLASSES = { + "BrinkmanStressKernel": (BrinkmanStressComponentKernel, 2027), + "BrinkmanletKernel": (BrinkmanletComponentKernel, 2027), + "ElasticityKernel": (ElasticityComponentKernel, 2027), + "Kernel": (ScalarKernel, 2027), + "StokesletKernel": (StokesletComponentKernel, 2027), + "StressletKernel": (StressletComponentKernel, 2027), +} + + +def __getattr__(name: str) -> Any: + result = _DEPRECATED_CLASSES.get(name) + + if result is not None: + cls, year = result + from warnings import warn + warn(f"'sumpy.kernel.{name}' is deprecated. " + f"Use 'sumpy.kernel.{cls.__name__}' instead. " + f"'sumpy.kernel.{name}' will continue to work until {year}.", + DeprecationWarning, stacklevel=2) + return cls + else: + raise AttributeError(name) + +# }}} + # vim: fdm=marker 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..2375bf4e 100644 --- a/sumpy/test/test_kernels.py +++ b/sumpy/test/test_kernels.py @@ -62,17 +62,17 @@ from sumpy.kernel import ( AxisTargetDerivative, BiharmonicKernel, - BrinkmanletKernel, - BrinkmanStressKernel, + BrinkmanletComponentKernel, + BrinkmanStressComponentKernel, DirectionalSourceDerivative, - ElasticityKernel, + ElasticityComponentKernel, HelmholtzKernel, - Kernel, LaplaceKernel, LineOfCompressionKernel, OneKernel, - StokesletKernel, - StressletKernel, + ScalarKernel, + StokesletComponentKernel, + StressletComponentKernel, YukawaKernel, ) from sumpy.test.geometries import make_ellipsoid, make_torus @@ -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() @@ -165,7 +165,7 @@ def test_p2e_multiple( extra_kwargs["k"] = 0.2 * (0.707 + 0.707j) else: extra_kwargs["k"] = 0.2 - if isinstance(base_knl, StokesletKernel): + if isinstance(base_knl, StokesletComponentKernel): extra_kwargs["mu"] = 0.2 source_kernels = [ @@ -310,7 +310,7 @@ def test_p2e2p( extra_kwargs["k"] = 0.2 * (0.707 + 0.707j) else: extra_kwargs["k"] = 0.2 - if isinstance(base_knl, StokesletKernel): + if isinstance(base_knl, StokesletComponentKernel): extra_kwargs["mu"] = 0.2 if with_source_derivative: @@ -478,13 +478,13 @@ def test_p2e2p( (HelmholtzKernel(3), LinearPDEConformingVolumeTaylorLocalExpansion, LinearPDEConformingVolumeTaylorMultipoleExpansion, False), (HelmholtzKernel(2), H2DLocalExpansion, H2DMultipoleExpansion, False), - (StokesletKernel(2, 0, 0), VolumeTaylorLocalExpansion, + (StokesletComponentKernel(2, 0, 0), VolumeTaylorLocalExpansion, VolumeTaylorMultipoleExpansion, False), - (StokesletKernel(2, 0, 0), LinearPDEConformingVolumeTaylorLocalExpansion, + (StokesletComponentKernel(2, 0, 0), LinearPDEConformingVolumeTaylorLocalExpansion, LinearPDEConformingVolumeTaylorMultipoleExpansion, False), - (BrinkmanletKernel(2, 0, 0), VolumeTaylorLocalExpansion, + (BrinkmanletComponentKernel(2, 0, 0), VolumeTaylorLocalExpansion, VolumeTaylorMultipoleExpansion, False), - (BrinkmanStressKernel(2, 0, 0, 0), VolumeTaylorLocalExpansion, + (BrinkmanStressComponentKernel(2, 0, 0, 0), VolumeTaylorLocalExpansion, VolumeTaylorMultipoleExpansion, False), ]) def test_translations( @@ -507,12 +507,12 @@ def test_translations( extra_kwargs["k"] = 0.05 elif isinstance(knl, YukawaKernel): extra_kwargs["lam"] = 0.05 - elif isinstance(knl, (StokesletKernel, StressletKernel)): + elif isinstance(knl, (StokesletComponentKernel, StressletComponentKernel)): extra_kwargs["mu"] = 0.05 - elif isinstance(knl, ElasticityKernel): + elif isinstance(knl, ElasticityComponentKernel): extra_kwargs["mu"] = 0.05 extra_kwargs["nu"] = 0.1 - elif isinstance(knl, (BrinkmanletKernel, BrinkmanStressKernel)): + elif isinstance(knl, (BrinkmanletComponentKernel, BrinkmanStressComponentKernel)): extra_kwargs["mu"] = 0.1 extra_kwargs["k"] = 0.1 @@ -554,7 +554,7 @@ def test_translations( del eval_offset if knl.dim == 2: - if isinstance(knl, BrinkmanStressKernel): # noqa: SIM108 + if isinstance(knl, BrinkmanStressComponentKernel): orders = [3, 4, 5] else: orders = [2, 3, 4] @@ -649,7 +649,7 @@ def test_m2m_and_l2l_exprs_simpler(base_knl, local_expn_class, mpole_expn_class, extra_kwargs["k"] = 0.2 * (0.707 + 0.707j) else: extra_kwargs["k"] = 0.2 - if isinstance(base_knl, StokesletKernel): + if isinstance(base_knl, StokesletComponentKernel): extra_kwargs["mu"] = 0.2 if with_source_derivative: @@ -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,30 +907,30 @@ 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__}" @pytest.mark.parametrize("knl", [ BiharmonicKernel(2), - BrinkmanletKernel(2, 0, 1), - BrinkmanletKernel(3, 0, 0, + BrinkmanletComponentKernel(2, 0, 1), + BrinkmanletComponentKernel(3, 0, 0, viscosity_mu_name="viscosity", darcy_impermeability_name="kappa"), - BrinkmanStressKernel(2, 0, 1, 0), - BrinkmanStressKernel(3, 0, 0, 1, + BrinkmanStressComponentKernel(2, 0, 1, 0), + BrinkmanStressComponentKernel(3, 0, 0, 1, viscosity_mu_name="viscosity", darcy_impermeability_name="kappa"), - ElasticityKernel(2, 0, 0), + ElasticityComponentKernel(2, 0, 0), HelmholtzKernel(3, helmholtz_k_name="kay"), LaplaceKernel(3), LineOfCompressionKernel(), OneKernel(2), - StokesletKernel(2, 0, 0), - StressletKernel(2, 0, 0, 0), + StokesletComponentKernel(2, 0, 0), + StressletComponentKernel(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..4412e1e9 100644 --- a/sumpy/test/test_misc.py +++ b/sumpy/test/test_misc.py @@ -58,17 +58,23 @@ ) from sumpy.kernel import ( BiharmonicKernel, - BrinkmanletKernel, - BrinkmanStressKernel, - ElasticityKernel, + BrinkmanletComponentKernel, + BrinkmanletSystemKernel, + BrinkmanStressComponentKernel, + BrinkmanStressSystemKernel, + ElasticityComponentKernel, + ElasticitySystemKernel, ExpressionKernel, HeatKernel, HelmholtzKernel, - Kernel, LaplaceKernel, LineOfCompressionKernel, - StokesletKernel, - StressletKernel, + ScalarKernel, + StokesletComponentKernel, + StokesletSystemKernel, + StressletComponentKernel, + StressletSystemKernel, + SystemKernel, 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() @@ -124,29 +130,29 @@ def nderivs(self): KernelInfo(LaplaceKernel(3)), KernelInfo(HelmholtzKernel(2), k=5), KernelInfo(HelmholtzKernel(3), k=5), - KernelInfo(StokesletKernel(2, 0, 1), mu=5), - KernelInfo(StokesletKernel(2, 1, 1), mu=5), - KernelInfo(StokesletKernel(3, 0, 1), mu=5), - KernelInfo(StokesletKernel(3, 1, 1), mu=5), - KernelInfo(StressletKernel(2, 0, 0, 0), mu=5), - KernelInfo(StressletKernel(2, 0, 0, 1), mu=5), - KernelInfo(StressletKernel(3, 0, 0, 0), mu=5), - KernelInfo(StressletKernel(3, 0, 0, 1), mu=5), - KernelInfo(StressletKernel(3, 0, 1, 2), mu=5), - KernelInfo(ElasticityKernel(2, 0, 1), mu=5, nu=0.2), - KernelInfo(ElasticityKernel(2, 0, 0), mu=5, nu=0.2), - KernelInfo(ElasticityKernel(3, 0, 1), mu=5, nu=0.2), - KernelInfo(ElasticityKernel(3, 0, 0), mu=5, nu=0.2), + KernelInfo(StokesletComponentKernel(2, 0, 1), mu=5), + KernelInfo(StokesletComponentKernel(2, 1, 1), mu=5), + KernelInfo(StokesletComponentKernel(3, 0, 1), mu=5), + KernelInfo(StokesletComponentKernel(3, 1, 1), mu=5), + KernelInfo(StressletComponentKernel(2, 0, 0, 0), mu=5), + KernelInfo(StressletComponentKernel(2, 0, 0, 1), mu=5), + KernelInfo(StressletComponentKernel(3, 0, 0, 0), mu=5), + KernelInfo(StressletComponentKernel(3, 0, 0, 1), mu=5), + KernelInfo(StressletComponentKernel(3, 0, 1, 2), mu=5), + KernelInfo(ElasticityComponentKernel(2, 0, 1), mu=5, nu=0.2), + KernelInfo(ElasticityComponentKernel(2, 0, 0), mu=5, nu=0.2), + KernelInfo(ElasticityComponentKernel(3, 0, 1), mu=5, nu=0.2), + KernelInfo(ElasticityComponentKernel(3, 0, 0), mu=5, nu=0.2), KernelInfo(LineOfCompressionKernel(3, 0), mu=5, nu=0.2), KernelInfo(LineOfCompressionKernel(3, 1), mu=5, nu=0.2), - KernelInfo(BrinkmanletKernel(2, 0, 1), mu=5, k=3), - KernelInfo(BrinkmanletKernel(2, 1, 1), mu=5, k=3), - KernelInfo(BrinkmanletKernel(3, 0, 1), mu=5, k=3), - KernelInfo(BrinkmanletKernel(3, 1, 1), mu=5, k=3), - KernelInfo(BrinkmanStressKernel(2, 0, 1, 0), mu=5, k=3), - KernelInfo(BrinkmanStressKernel(2, 0, 1, 1), mu=5, k=3), - KernelInfo(BrinkmanStressKernel(3, 0, 1, 0), mu=5, k=3), - KernelInfo(BrinkmanStressKernel(3, 0, 1, 2), mu=5, k=3), + KernelInfo(BrinkmanletComponentKernel(2, 0, 1), mu=5, k=3), + KernelInfo(BrinkmanletComponentKernel(2, 1, 1), mu=5, k=3), + KernelInfo(BrinkmanletComponentKernel(3, 0, 1), mu=5, k=3), + KernelInfo(BrinkmanletComponentKernel(3, 1, 1), mu=5, k=3), + KernelInfo(BrinkmanStressComponentKernel(2, 0, 1, 0), mu=5, k=3), + KernelInfo(BrinkmanStressComponentKernel(2, 0, 1, 1), mu=5, k=3), + KernelInfo(BrinkmanStressComponentKernel(3, 0, 1, 0), mu=5, k=3), + KernelInfo(BrinkmanStressComponentKernel(3, 0, 1, 2), mu=5, k=3), KernelInfo(HeatKernel(1), alpha=0.1), KernelInfo(HeatKernel(2), alpha=0.1), KernelInfo(HeatKernel(3), alpha=0.1), @@ -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) @@ -561,9 +567,9 @@ def test_as_scalar_pde_brinkman(): def test_elasticity_pickle(): from pickle import dumps, loads - stokes_knl = StokesletKernel( + stokes_knl = StokesletComponentKernel( 3, 0, 1, viscosity_mu_name="mu1") - elasticity_knl = ElasticityKernel( + elasticity_knl = ElasticityComponentKernel( 3, 0, 1, viscosity_mu_name="mu1", poisson_ratio_name="nu1") elasticity_helper_knl = LineOfCompressionKernel( 3, 0, viscosity_mu_name="mu1", poisson_ratio_name="nu1") @@ -790,6 +796,83 @@ def test_get_storage_index(order, knl, compressed): # }}} +# {{{ test_system_kernel_components + +def _get_scalar_cls(knl: SystemKernel) -> type[ScalarKernel]: + if isinstance(knl, ElasticitySystemKernel): + return ElasticityComponentKernel + elif isinstance(knl, StokesletSystemKernel): + return StokesletComponentKernel + elif isinstance(knl, StressletSystemKernel): + return StressletComponentKernel + elif isinstance(knl, BrinkmanletSystemKernel): + return BrinkmanletComponentKernel + elif isinstance(knl, BrinkmanStressSystemKernel): + return BrinkmanStressComponentKernel + 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[SystemKernel]) -> 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] == 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] == 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[SystemKernel]) -> 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 7d2847cc..5ee91357 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,