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 de6f3eeb..821a3f0c 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -44,7 +44,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 +53,6 @@ ExprDerivativeTaker, diff_derivative_coeff_dict, ) -from sumpy.symbolic import SpatialConstant, pymbolic_real_norm_2 if TYPE_CHECKING: @@ -69,7 +68,11 @@ ---------------- .. autoclass:: KernelArgument -.. autoclass:: Kernel +.. autoclass:: GenericKernel + :show-inheritance: +.. autoclass:: ScalarKernel + :show-inheritance: +.. autoclass:: SystemKernel :show-inheritance: Symbolic kernels @@ -79,6 +82,10 @@ :show-inheritance: :members: mapper_method +.. autoclass:: OneKernel + :show-inheritance: + :members: mapper_method + PDE kernels ----------- @@ -195,8 +202,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 @@ -205,32 +229,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: @@ -246,19 +257,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. @@ -274,32 +348,6 @@ def get_code_transformer(self) -> Callable[[Expression], Expression]: """ 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, @@ -313,6 +361,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 @@ -340,7 +401,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 @@ -370,49 +431,46 @@ 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 + @abstractmethod + def shape(self) -> tuple[int, ...]: + pass - def get_source_args(self) -> Sequence[KernelArgument]: + @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 @@ -443,24 +501,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( @@ -498,6 +553,51 @@ 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: obj_array.ObjectArrayND[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 shape(self) -> tuple[int, ...]: + return self.components.shape + + @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.components.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): @@ -513,11 +613,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: @@ -566,7 +666,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, @@ -640,17 +740,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: @@ -686,9 +786,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( @@ -729,7 +828,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) @@ -740,7 +839,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) @@ -783,9 +882,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( @@ -855,11 +953,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: @@ -958,10 +1056,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: @@ -1057,7 +1155,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 @@ -1157,12 +1255,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 @@ -1249,11 +1347,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 @@ -1374,10 +1472,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 @@ -1472,15 +1570,85 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator: # }}} +# {{{ vector PDE kernels + + +@dataclass(frozen=True, repr=False) +class ElasticitySystemKernel(ExpressionSystemKernel): + mapper_method: ClassVar[str] = "map_elasticity_system_kernel" + + viscosity_mu_name: str + poisson_ratio_name: str + + def __init__(self, + dim: int, + viscosity_mu_name: str = "mu", + poisson_ratio_name: str = "nu") -> None: + components = np.zeros((dim, dim), dtype=object) + + 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=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, 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)) + + +# }}} + # {{{ 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 @@ -1489,7 +1657,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 @@ -1539,7 +1707,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__}'") @@ -1557,7 +1725,7 @@ def get_derivative_taker(self, # {{{ derivatives class DerivativeBase(KernelWrapper, ABC): - """Bases: :class:`Kernel` + """Bases: :class:`ScalarKernel` .. autoattribute:: inner_kernel .. automethod:: replace_inner_kernel @@ -1568,11 +1736,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. """ @@ -1587,7 +1755,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) @@ -1609,12 +1777,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) @@ -1629,7 +1797,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) @@ -1667,12 +1835,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) @@ -1720,7 +1888,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" @@ -1733,13 +1903,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) @@ -1763,8 +1933,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) @@ -1799,7 +1968,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. @@ -1813,7 +1982,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) @@ -1870,11 +2039,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) # }}} @@ -1889,17 +2058,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) @@ -1929,42 +2098,32 @@ 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 - - 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 + + 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) @@ -1973,7 +2132,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) @@ -1981,7 +2140,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) @@ -1994,7 +2153,7 @@ class SourceDerivativeRemover(AxisSourceDerivativeRemover): @override def map_directional_source_derivative( - self, kernel: DirectionalSourceDerivative) -> Kernel: + self, kernel: DirectionalSourceDerivative) -> ScalarKernel: return self.rec(kernel.inner_kernel) @@ -2002,7 +2161,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 297bfe09..96e4a2bb 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() @@ -892,7 +892,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, @@ -955,7 +955,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__}" @@ -978,7 +978,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 c6c8245e..f2cdbcab 100644 --- a/sumpy/test/test_misc.py +++ b/sumpy/test/test_misc.py @@ -63,9 +63,9 @@ ElasticityKernel, ExpressionKernel, HelmholtzKernel, - Kernel, LaplaceKernel, LineOfCompressionKernel, + ScalarKernel, StokesletKernel, StressletKernel, YukawaKernel, @@ -214,7 +214,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) 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 3ccda9c8..66bc2c44 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,