from __future__ import annotations

__copyright__ = "Copyright (C) 2012 Andreas Kloeckner"

__license__ = """

Permission is hereby granted, free of charge, to any person obtaining a copy

of this software and associated documentation files (the "Software"), to deal

in the Software without restriction, including without limitation the rights

to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

copies of the Software, and to permit persons to whom the Software is

furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in

all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

THE SOFTWARE.

"""

from abc import ABC, abstractmethod

from collections import defaultdict

from dataclasses import dataclass

from typing import (

TYPE_CHECKING,

ClassVar,

Generic,

Literal,

TypeVar,

cast,

overload,

)

import numpy as np

from typing_extensions import Self, override

import loopy as lp

import pymbolic.primitives as prim

from pymbolic import Expression, var

from pymbolic.mapper import CSECachingMapperMixin, IdentityMapper

from pymbolic.primitives import make_sym_vector

from pytools import memoize_method

import sumpy.symbolic as sym

from sumpy.derivative_taker import (

DerivativeCoeffDict,

DifferentiatedExprDerivativeTaker,

ExprDerivativeTaker,

diff_derivative_coeff_dict,

)

from sumpy.symbolic import SpatialConstant, pymbolic_real_norm_2

if TYPE_CHECKING:

from collections.abc import Callable, Iterable, Sequence

import sympy as sp

from sumpy.assignment_collection import SymbolicAssignmentCollection

from sumpy.expansion.diff_op import LinearPDESystemOperator

__doc__ = """

Kernel interface

----------------

.. autoclass:: KernelArgument

.. autoclass:: Kernel

:show-inheritance:

Symbolic kernels

----------------

.. autoclass:: ExpressionKernel

:show-inheritance:

:members: mapper_method

PDE kernels

-----------

.. autoclass:: LaplaceKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: BiharmonicKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: HelmholtzKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: YukawaKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: StokesletKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: StressletKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: ElasticityKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: LineOfCompressionKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: BrinkmanletKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: BrinkmanStressKernel

:show-inheritance:

:members: mapper_method

.. autoclass:: HeatKernel

:show-inheritance:

:members: mapper_method

.. [Pozrikidis1992] C. Pozrikidis,

*Boundary Integral and Singularity Methods for Linearized Viscous Flow*,

Cambridge University Press, 1992.

.. [Hsiao2008] G. C. Hsiao, W. L. Wendland,

*Boundary Integral Equations*,

Springer, 2008.

.. [Kress2013] R. Kress,

*Linear Integral Equations*,

Springer Science & Business Media, 2013.

Derivatives

-----------

These objects *wrap* other kernels and take derivatives

of them in the process.

.. autoclass:: DerivativeBase

.. autoclass:: AxisTargetDerivative

:show-inheritance:

:undoc-members:

:members: mapper_method,target_array_name

.. autoclass:: AxisSourceDerivative

:show-inheritance:

:members: mapper_method

.. autoclass:: DirectionalDerivative

:show-inheritance:

:members: directional_kind

.. autoclass:: DirectionalSourceDerivative

:show-inheritance:

:members: mapper_method,directional_kind

Transforming kernels

--------------------

.. autoclass:: TargetPointMultiplier

:undoc-members:

:members: mapper_method,target_array_name

.. autoclass:: ResultT

.. autoclass:: KernelMapper

:show-inheritance:

.. autoclass:: KernelCombineMapper

:show-inheritance:

.. autoclass:: KernelIdentityMapper

:show-inheritance:

.. autoclass:: AxisSourceDerivativeRemover

:show-inheritance:

.. autoclass:: AxisTargetDerivativeRemover

:show-inheritance:

.. autoclass:: SourceDerivativeRemover

:show-inheritance:

.. autoclass:: TargetDerivativeRemover

:show-inheritance:

.. autoclass:: TargetTransformationRemover

:show-inheritance:

.. autoclass:: DerivativeCounter

:show-inheritance:

"""

@dataclass(frozen=True)

class KernelArgument:

"""

.. autoattribute:: loopy_arg

"""

loopy_arg: lp.KernelArgument

"""A :class:`loopy.KernelArgument` instance describing the type, name, and

other features of this kernel argument when passed to a generated piece of

code.

"""

@property

def name(self) -> str:

return self.loopy_arg.name

# {{{ basic kernel interface

@dataclass(frozen=True)

class Kernel(ABC):

"""Basic kernel interface.

.. autoattribute:: mapper_method

.. autoattribute:: is_translation_invariant

.. 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_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."""

@override

def __repr__(self) -> str:

from dataclasses import fields

args: list[str] = []

for f in fields(self):

value = getattr(self, f.name)

if isinstance(value, prim.ExpressionNode):

args.append(f"{f.name}={value}")

else:

args.append(f"{f.name}={value!r}")

return f"{type(self).__name__}({', '.join(args)})"

def get_base_kernel(self) -> Kernel:

"""

:returns: the kernel being wrapped by this one, or else *self*.

"""

return self

def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel:

"""

:returns: the base kernel being wrapped by this one, or else

*new_base_kernel*.

"""

return new_base_kernel

def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit:

"""Apply some changes (such as registering function manglers) to the kernel.

:returns: a new :mod:`loopy` kernel with the applied changes.

"""

return loopy_knl

def get_code_transformer(self) -> Callable[[Expression], Expression]:

"""

:returns: a function to postprocess the :mod:`pymbolic` expression

generated from the result of :meth:`get_expression` on the way to

code generation.

"""

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,

dvec: sp.Matrix,

rscale: sym.Expr,

sac: SymbolicAssignmentCollection | None,

) -> ExprDerivativeTaker:

"""

:returns: an :class:`~sumpy.derivative_taker.ExprDerivativeTaker` instance

that supports taking derivatives of the base kernel with respect to

*dvec*.

"""

@overload

def postprocess_at_source(

self, expr: sym.Expr, avec: sym.Matrix

) -> sym.Expr: ...

@overload

def postprocess_at_source(

self, expr: ExprDerivativeTaker, avec: sym.Matrix

) -> DifferentiatedExprDerivativeTaker: ...

def postprocess_at_source(

self, expr: sym.Expr | ExprDerivativeTaker, avec: sym.Matrix,

) -> sym.Expr | DifferentiatedExprDerivativeTaker:

"""Transform a kernel evaluation or expansion expression in a place

where the vector :math:`a` (something - source) is known. ("something" may be

an expansion center or a target)

The typical use of this function is to apply source-variable

derivatives to the kernel.

"""

expr_dict: DerivativeCoeffDict = {(0,)*self.dim: 1}

expr_dict = self.get_derivative_coeff_dict_at_source(expr_dict)

if isinstance(expr, ExprDerivativeTaker):

return DifferentiatedExprDerivativeTaker(expr, expr_dict)

result = 0

for mi, coeff in expr_dict.items():

result += coeff * self._diff(expr, avec, mi)

assert isinstance(result, sym.Expr)

return result

@overload

def postprocess_at_target(

self, expr: sym.Expr, bvec: sym.Matrix,

) -> sym.Expr: ...

@overload

def postprocess_at_target(

self, expr: ExprDerivativeTaker | DifferentiatedExprDerivativeTaker,

bvec: sym.Matrix,

) -> DifferentiatedExprDerivativeTaker: ...

def postprocess_at_target(self,

expr:

sym.Expr | ExprDerivativeTaker | DifferentiatedExprDerivativeTaker,

bvec: sym.Matrix,

) -> sym.Expr | DifferentiatedExprDerivativeTaker:

"""Transform a kernel evaluation or expansion expression in a place

where the vector :math:`b` (target - something) is known. ("something" may

be an expansion center or a target.)

The typical use of this function is to apply target-variable

derivatives to the kernel.

"""

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`.

This function is meant to be overridden by child classes where necessary.

"""

return expr_dict

@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.

"""

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 []

# }}}

@dataclass(frozen=True, repr=False)

class ExpressionKernel(Kernel, ABC):

r"""

.. autoattribute:: expression

.. autoattribute:: global_scaling_const

"""

mapper_method: ClassVar[str] = "map_expression_kernel"

expression: Expression

"""A :mod:`pymbolic` expression depending on variables *d_1* through *d_N*

where *N* equals *dim*. These variables match what is returned from

:func:`pymbolic.primitives.make_sym_vector` with argument `"d"`. Any

variable that is not *d* or a :class:`~sumpy.symbolic.SpatialConstant` will

be viewed as potentially spatially varying.

"""

global_scaling_const: Expression

r"""A constant :mod:`pymbolic` expression for the global scaling 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.

"""

@override

def __str__(self) -> str:

return f"ExprKnl{self.dim}D"

@override

def get_expression(self, dist_vec: sym.Matrix) -> sym.Expr:

from sumpy.symbolic import PymbolicToSympyMapperWithSymbols

expr = 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

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)

@override

def get_derivative_taker(

self,

dvec: sym.Matrix,

rscale: sym.Expr,

sac: SymbolicAssignmentCollection | None,

) -> ExprDerivativeTaker:

return ExprDerivativeTaker(self.get_expression(dvec), dvec, rscale, sac)

@override

def get_pde_as_diff_op(self) -> LinearPDESystemOperator:

raise NotImplementedError

class OneKernel(ExpressionKernel):

def __init__(self, dim: int):

super().__init__(

dim=dim,

expression=1,

global_scaling_const=1,

)

@override

def __reduce__(self) -> tuple[object, ...]:

return (self.__class__, (self.dim,))

@property

@override

def is_complex_valued(self) -> bool:

return False

one_kernel_2d = OneKernel(2)

one_kernel_3d = OneKernel(3)

# {{{ PDE kernels

class LaplaceKernel(ExpressionKernel):

r"""A kernel for the Laplace equation (see e.g. Theorem 6.2 from [Kress2013]_).

.. math::

\Delta K(\mathbf{x}, \mathbf{y}) = \delta(\mathbf{x} - \mathbf{y}).

"""

mapper_method: ClassVar[str] = "map_laplace_kernel"

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))

expr = var("log")(r)

scaling = 1/(-2*var("pi"))

elif dim == 3:

r = pymbolic_real_norm_2(make_sym_vector("d", dim))

expr = 1/r

scaling = 1/(4*var("pi"))

else:

raise NotImplementedError(f"unsupported dimension: '{dim}'")

super().__init__(dim, expression=expr, global_scaling_const=scaling)

@override

def __reduce__(self) -> tuple[object, ...]:

return (self.__class__, (self.dim,))

@property

@override

def is_complex_valued(self) -> bool:

return False

@override

def __str__(self) -> str:

return f"LapKnl{self.dim}D"

@override

def get_derivative_taker(

self,

dvec: sym.Matrix,

rscale: sym.Expr,

sac: SymbolicAssignmentCollection | None,

) -> ExprDerivativeTaker:

from sumpy.derivative_taker import LaplaceDerivativeTaker

return LaplaceDerivativeTaker(self.get_expression(dvec), dvec, rscale, sac)

@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)

return laplacian(w)

class BiharmonicKernel(ExpressionKernel):

r"""A kernel for the biharmonic equation.

.. math::

\Delta^2 K(\mathbf{x}, \mathbf{y}) = \delta(\mathbf{x} - \mathbf{y}).

"""

mapper_method: ClassVar[str] = "map_biharmonic_kernel"

def __init__(self, dim: int) -> None:

r = 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,

# Department of Computer Science, Yale University, 1990.

expr = r**2 * var("log")(r)

scaling = 1/(8*var("pi"))

elif dim == 3:

# Ref: Jiang, Shidong, Bo Ren, Paul Tsuji, and Lexing Ying.

# "Second kind integral equations for the first kind Dirichlet problem

# of the biharmonic equation in three dimensions."

# Journal of Computational Physics 230, no. 19 (2011): 7488-7501.

expr = r

scaling = -1/(8*var("pi"))

else:

raise NotImplementedError(f"unsupported dimension: '{dim}'")

super().__init__(dim, expression=expr, global_scaling_const=scaling)

@override

def __reduce__(self) -> tuple[object, ...]:

return (self.__class__, (self.dim,))

@property

@override

def is_complex_valued(self) -> bool:

return False

@override

def __str__(self) -> str:

return f"BiharmKnl{self.dim}D"

@override

def get_derivative_taker(

self,

dvec: sp.Matrix,

rscale: sym.Expr,

sac: SymbolicAssignmentCollection | None,

) -> ExprDerivativeTaker:

from sumpy.derivative_taker import RadialDerivativeTaker

return RadialDerivativeTaker(self.get_expression(dvec), dvec, rscale, sac)

@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)

return laplacian(laplacian(w))

@dataclass(frozen=True, repr=False)

class HelmholtzKernel(ExpressionKernel):

r"""A kernel for the Helmholtz equation (see e.g. Example 12.14 in [Kress2013]_).

.. math::

\Delta K(\mathbf{x}, \mathbf{y}) + k^2 K(\mathbf{x}, \mathbf{y})

= \delta(\mathbf{x} - \mathbf{y}).

.. autoattribute:: helmholtz_k_name

.. autoattribute:: allow_evanescent

"""

mapper_method: ClassVar[str] = "map_helmholtz_kernel"

helmholtz_k_name: str

"""The argument name to use for the Helmholtz parameter when generating

functions to evaluate this kernel.

"""

allow_evanescent: bool

def __init__(self,

dim: int,

helmholtz_k_name: str = "k",

allow_evanescent: bool = False) -> None:

k = 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))

expr = var("hankel_1")(0, k*r)

scaling = var("I")/4

elif dim == 3:

r = pymbolic_real_norm_2(make_sym_vector("d", dim))

expr = var("exp")(var("I")*k*r)/r

scaling = 1/(4*var("pi"))

else:

raise NotImplementedError(f"unsupported dimension: '{dim}'")

super().__init__(dim, expression=expr, global_scaling_const=scaling)

object.__setattr__(self, "helmholtz_k_name", helmholtz_k_name)

object.__setattr__(self, "allow_evanescent", allow_evanescent)

@override

def __reduce__(self) -> tuple[object, ...]:

return (

self.__class__,

(self.dim, self.helmholtz_k_name, self.allow_evanescent),

)

@property

@override

def is_complex_valued(self) -> bool:

return True

@override

def __str__(self) -> str:

return f"HelmKnl{self.dim}D({self.helmholtz_k_name})"

@override

def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit:

from sumpy.codegen import register_bessel_callables

return register_bessel_callables(loopy_knl)

@override

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),

)]

@override

def get_derivative_taker(

self,

dvec: sym.Matrix,

rscale: sym.Expr,

sac: SymbolicAssignmentCollection | None,

) -> ExprDerivativeTaker:

from sumpy.derivative_taker import HelmholtzDerivativeTaker

return HelmholtzDerivativeTaker(self.get_expression(dvec), dvec, rscale, sac)

@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.helmholtz_k_name)

return laplacian(w) + k**2 * w

@dataclass(frozen=True, repr=False)

class YukawaKernel(ExpressionKernel):

r"""A kernel for the Yukawa equation.

.. math::

\Delta K(\mathbf{x}, \mathbf{y}) - \lambda^2 K(\mathbf{x}, \mathbf{y})

= \delta(\mathbf{x} - \mathbf{y}).

.. autoattribute:: yukawa_lambda_name

"""

mapper_method: ClassVar[str] = "map_yukawa_kernel"

yukawa_lambda_name: str

"""The argument name to use for the Yukawa parameter when generating

functions to evaluate this kernel.

"""

def __init__(self, dim: int, yukawa_lambda_name: str = "lam") -> None:

lam = 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)

# where K is a modified Bessel function of the second kind.

#

# [1] https://en.wikipedia.org/wiki/Green%27s_function

# [2] https://dlmf.nist.gov/10.27#E8

# [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))

if dim == 2:

# NOTE: transform K(0, lam r) into a Hankel function using [2]

expr = var("hankel_1")(0, var("I")*lam*r)

scaling_for_K0 = var("pi")/2*var("I") # noqa: N806

scaling = 1/(2*var("pi")) * scaling_for_K0

elif dim == 3:

# NOTE: to get the expression, we do the following and simplify

# 1. express K(1/2, lam r) as a modified spherical Bessel function

# k(0, lam r) using [3] and use expression for k(0, lam r) from [4]

# 2. or use (AS 10.2.17) directly

expr = var("exp")(-lam*r) / r

scaling = 1/(4 * var("pi"))

else:

raise NotImplementedError(f"unsupported dimension: '{dim}'")

super().__init__(dim, expression=expr, global_scaling_const=scaling)

object.__setattr__(self, "yukawa_lambda_name", yukawa_lambda_name)

@override

def __reduce__(self) -> tuple[object, ...]:

return (self.__class__, (self.dim, self.yukawa_lambda_name))

@property

@override

def is_complex_valued(self) -> bool:

# FIXME: 2D uses Hankel functions (complex-valued); 3D uses real exp

return True

@override

def __str__(self) -> str:

return f"YukKnl{self.dim}D({self.yukawa_lambda_name})"

@override

def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit:

from sumpy.codegen import register_bessel_callables

return register_bessel_callables(loopy_knl)

@override

def get_args(self) -> Sequence[KernelArgument]:

return [

KernelArgument(

loopy_arg=lp.ValueArg(self.yukawa_lambda_name, np.float64),

)]

@override

def get_derivative_taker(

self,

dvec: sym.Matrix,

rscale: sym.Expr,

sac: SymbolicAssignmentCollection | None,

) -> ExprDerivativeTaker:

from sumpy.derivative_taker import HelmholtzDerivativeTaker

return HelmholtzDerivativeTaker(self.get_expression(dvec), dvec, rscale, sac)

@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)

lam = sym.Symbol(self.yukawa_lambda_name)

return laplacian(w) - lam**2 * w

@dataclass(frozen=True, repr=False)

class ElasticityKernel(ExpressionKernel):

r"""A kernel for the linear elasticity (Navier-Cauchy) equations

(see e.g. Section 2.2 in [Hsiao2008]_).

.. math::

\mu \Delta K_{ij}(\mathbf{x}, \mathbf{y})

+ \frac{\mu}{1 - 2 \nu} \nabla (\nabla \cdot K_{ij}(\mathbf{x}, \mathbf{y}))

= \delta_{ij} \delta(\mathbf{x} - \mathbf{y}).

.. autoattribute:: icomp

.. autoattribute:: jcomp

.. autoattribute:: viscosity_mu_name

.. autoattribute:: poisson_ratio_name

"""

mapper_method: ClassVar[str] = "map_elasticity_kernel"

icomp: int

"""Component index for the kernel."""

jcomp: int

"""Component index for the kernel."""

viscosity_mu_name: str

r"""The argument name to use for the dynamic viscosity :math:`\mu` when

generating functions to evaluate this kernel.

"""

poisson_ratio_name: str

r"""The argument name to use for Poisson's ratio :math:`\nu` when generating

functions to evaluate this kernel.

"""

def __init__(self,

dim: int,

icomp: int,

jcomp: int,

viscosity_mu_name: str = "mu",

poisson_ratio_name: str = "nu") -> None:

if not isinstance(viscosity_mu_name, str):

raise TypeError(

f"'viscosity_mu_name' is not a str: {type(viscosity_mu_name)}"

)

if not isinstance(poisson_ratio_name, str):

raise TypeError(

f"'poisson_ratio_name' is not a str: {type(poisson_ratio_name)}"

)

mu = SpatialConstant(viscosity_mu_name)

nu = SpatialConstant(poisson_ratio_name)

d = make_sym_vector("d", dim)

r = pymbolic_real_norm_2(d)

delta_ij = 1 if icomp == jcomp else 0

if dim == 2:

# See (Berger and Karageorghis 2001)

expr = -var("log")(r)*(3 - 4 * nu)*delta_ij + d[icomp]*d[jcomp]/r**2

scaling = -1/(8*var("pi")*(1 - nu)*mu)

elif dim == 3:

# Kelvin solution

expr = (1/r)*(3 - 4*nu)*delta_ij + d[icomp]*d[jcomp]/r**3

scaling = -1/(16*var("pi")*(1 - nu)*mu)

else:

raise NotImplementedError(f"unsupported dimension: '{dim}'")

super().__init__(dim, expression=expr, global_scaling_const=scaling)

object.__setattr__(self, "icomp", icomp)

object.__setattr__(self, "jcomp", jcomp)

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_{self.icomp}{self.jcomp}"

f"({self.viscosity_mu_name}, {self.poisson_ratio_name})")

@override

def __reduce__(self):

return (

type(self),

(self.dim, self.icomp, self.jcomp,

self.viscosity_mu_name,

self.poisson_ratio_name))

@property

@override

def is_complex_valued(self) -> bool:

return False

@memoize_method

@override

def get_args(self) -> Sequence[KernelArgument]:

return [

KernelArgument(loopy_arg=lp.ValueArg(self.viscosity_mu_name, np.float64)),

KernelArgument(loopy_arg=lp.ValueArg(self.poisson_ratio_name, np.float64)),

]

@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)

return laplacian(laplacian(w))

@dataclass(frozen=True, repr=False)

class StokesletKernel(ExpressionKernel):

r"""A kernel for the Stokes equations (see e.g. Chapter 2 in [Pozrikidis1992]_).

.. math::

\begin{cases}

-\mu \Delta K_{ij}(\mathbf{x}, \mathbf{y})

+ \nabla_i P_j(\mathbf{x}, \mathbf{y})

= \delta_{ij} \delta(\mathbf{x} - \mathbf{y}), \\

\nabla_i K_{ij}(\mathbf{x}, \mathbf{y}) = 0, \\

\end{cases}

where pressure kernel :math:`P_j = \partial_j K` is the derivative of the

Laplace kernel. This kernel is often called the Stokeslet or the Oseen-Burgers

tensor and it represents the velocity field.

.. autoattribute:: icomp

.. autoattribute:: jcomp

.. autoattribute:: viscosity_mu_name

"""

mapper_method: ClassVar[str] = "map_stokeslet_kernel"

icomp: int

"""Component index for the kernel."""

jcomp: int

"""Component index for the kernel."""

viscosity_mu_name: str

r"""The argument name to use for the dynamic viscosity :math:`\mu` when

generating functions to evaluate this kernel.

"""

def __init__(self,

dim: int,

icomp: int,

jcomp: int,

viscosity_mu_name: str = "mu") -> None:

if not isinstance(viscosity_mu_name, str):

raise TypeError(

f"'viscosity_mu_name' is not a str: {type(viscosity_mu_name)}"

)

mu = SpatialConstant(viscosity_mu_name)

d = make_sym_vector("d", dim)

r = pymbolic_real_norm_2(d)

delta_ij = 1 if icomp == jcomp else 0

if dim == 2:

expr = -var("log")(r)*delta_ij + d[icomp]*d[jcomp]/r**2

scaling = -1/(4*var("pi")*mu)

elif dim == 3:

expr = (1/r)*delta_ij + d[icomp]*d[jcomp]/r**3

scaling = -1/(8*var("pi")*mu)

else:

raise NotImplementedError(f"unsupported dimension: '{dim}'")

super().__init__(dim, expression=expr, global_scaling_const=scaling)

object.__setattr__(self, "icomp", icomp)

object.__setattr__(self, "jcomp", jcomp)

object.__setattr__(self, "viscosity_mu_name", viscosity_mu_name)

@override

def __str__(self) -> str:

return (

f"StokesletKnl{self.dim}D_{self.icomp}{self.jcomp}"

f"({self.viscosity_mu_name})")

@override

def __reduce__(self):

return (

type(self),

(self.dim, self.icomp, self.jcomp, self.viscosity_mu_name))

@property

@override

def is_complex_valued(self) -> bool:

return False