Add heat kernel

.github/workflows/ci.yml

@@ -68,6 +68,7 @@ jobs:
             run: |
                 grep -v symengine .test-conda-env-py3.yml > .test-conda-env.yml
                 CONDA_ENVIRONMENT=.test-conda-env.yml
+                export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
                 curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
                 . ./build-and-test-py-project-within-miniconda.sh
 
@@ -89,6 +90,7 @@ jobs:
         -   uses: actions/checkout@v6
         -   name: "Main Script"
             run: |
+                export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
                 curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
                 . ./build-and-test-py-project-within-miniconda.sh
 

sumpy/e2e.py

@@ -320,9 +320,14 @@
                 [sym.Symbol(f"data{i}")
             for i in range(m2l_translation_classes_dependent_ndata)]
 
-        ncoeff_src = len(self.src_expansion)
+        m2l_translation = self.tgt_expansion.m2l_translation
+        if m2l_translation.use_preprocessing:
+            ncoeff_src = m2l_translation.preprocess_multipole_nexprs(
+                self.tgt_expansion, self.src_expansion)
+        else:
+            ncoeff_src = len(self.src_expansion)
 
         src_coeff_exprs = [sym.Symbol(f"src_coeffs{i}") for i in range(ncoeff_src)]
 
         from sumpy.assignment_collection import SymbolicAssignmentCollection
         sac = SymbolicAssignmentCollection()

sumpy/kernel.py

@@ -106,6 +106,9 @@
 .. autoclass:: LineOfCompressionKernel
     :show-inheritance:
     :members: mapper_method
+.. autoclass:: HeatKernel
+    :show-inheritance:
+    :members: mapper_method
 
 Derivatives
 -----------
@@ -1119,6 +1122,76 @@
         return laplacian(w)
 
 
+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.
+
+    .. 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.
+    """
+    heat_alpha_name: str
+
+    mapper_method: ClassVar[str] = "map_heat_kernel"
+
+    def __init__(self, spatial_dims: int, heat_alpha_name: str = "alpha"):
+        dim = spatial_dims + 1
+        d = make_sym_vector("d", dim)
+        t = d[-1]
+        r = pymbolic_real_norm_2(d[:-1])
+        alpha = SpatialConstant(heat_alpha_name)
+        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,
+                )
+
+        self.heat_alpha_name = heat_alpha_name
+
+    @property
+    @override
+    def is_complex_valued(self) -> bool:
+        return False
+
+    def update_persistent_hash(self, key_hash, key_builder):
+        key_hash.update(type(self).__name__.encode("utf8"))
+        key_builder.rec(key_hash, (self.dim, self.heat_alpha_name))
+
+    @override
+    def __repr__(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),
+                    )]
+
+    def get_derivative_taker(self, dvec, rscale, sac):
+        """Return a :class:`sumpy.derivative_taker.ExprDerivativeTaker` instance
+        that supports taking derivatives of the base kernel with respect to dvec.
+        """
+        from sumpy.derivative_taker import ExprDerivativeTaker
+        return ExprDerivativeTaker(self.get_expression(dvec), dvec, rscale,
+                sac)
+
+    @override
+    def get_pde_as_diff_op(self):
+        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
+
+
 # }}}
 
 
@@ -1590,6 +1663,7 @@
     map_elasticity_kernel = map_expression_kernel
     map_line_of_compression_kernel = map_expression_kernel
     map_stresslet_kernel = map_expression_kernel
+    map_heat_kernel = map_expression_kernel
 
     def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> Kernel:
         return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel))
@@ -1662,6 +1736,7 @@
     map_yukawa_kernel = map_expression_kernel
     map_line_of_compression_kernel = map_expression_kernel
     map_stresslet_kernel = map_expression_kernel
+    map_heat_kernel = map_expression_kernel
 
     @override
     def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> int: