Add heat kernel

.github/workflows/ci.yml

@@ -59,6 +59,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
 
@@ -69,6 +70,7 @@ jobs:
         -   uses: actions/checkout@v3
         -   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
 

.gitlab-ci.yml

@@ -27,6 +27,7 @@ Pytest POCL:
   - export PY_EXE=python3
   - export PYOPENCL_TEST=portable:pthread
   - export EXTRA_INSTALL="pybind11 numpy mako mpi4py"
+  - 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.sh
   - ". ./build-and-test-py-project.sh"
   tags:
@@ -44,6 +45,7 @@ Pytest Titan V:
   - py_version=3
   - export PYOPENCL_TEST=nvi:titan
   - EXTRA_INSTALL="pybind11 numpy mako mpi4py"
+  - 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.sh
   - ". ./build-and-test-py-project.sh"
   tags:
@@ -63,6 +65,7 @@ Pytest Conda:
   - export SUMPY_NO_CACHE=1
   - export SUMPY_FORCE_SYMBOLIC_BACKEND=symengine
   - export PYOPENCL_TEST=portable:pthread
+  - 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"
   tags:
@@ -80,6 +83,7 @@ Pytest POCL Titan V:
   - export SUMPY_NO_CACHE=1
   - export SUMPY_FORCE_SYMBOLIC_BACKEND=symengine
   - export PYOPENCL_TEST=portable:titan
+  - 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"
   tags:

examples/heat-toys.py

@@ -0,0 +1,80 @@
+import numpy as np
+
+import pyopencl as cl
+
+import sumpy.toys as t
+from sumpy.visualization import FieldPlotter
+from sumpy.kernel import (      # noqa: F401
+        YukawaKernel,
+        HeatKernel,
+        HelmholtzKernel,
+        LaplaceKernel)
+
+try:
+    import matplotlib.pyplot as plt
+    USE_MATPLOTLIB = True
+except ImportError:
+    USE_MATPLOTLIB = False
+
+
+def main():
+    from sumpy.array_context import PyOpenCLArrayContext
+    ctx = cl.create_some_context()
+    queue = cl.CommandQueue(ctx)
+    actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
+
+    tctx = t.ToyContext(
+            actx.context,
+            # LaplaceKernel(2),
+            #YukawaKernel(2), extra_kernel_kwargs={"lam": 5},
+            # HelmholtzKernel(2), extra_kernel_kwargs={"k": 0.3},
+            HeatKernel(1), extra_kernel_kwargs={"alpha": 1},
+            )
+
+    src_size = 1
+    pt_src = t.PointSources(
+            tctx,
+            np.array([
+                src_size*(np.random.rand(50) - 0.5),
+                np.zeros(50)]),
+            np.random.randn(50))
+
+    fp = FieldPlotter([0, 0.5], extent=np.array([8, 1]))
+
+    if 0 and USE_MATPLOTLIB:
+        t.logplot(fp, pt_src, cmap="jet", aspect=8)
+        plt.colorbar()
+        plt.show()
+
+
+    p = 5
+    mexp = t.multipole_expand(pt_src, [0, 0], p)
+    diff = mexp - pt_src
+
+    x, t = fp.points
+
+    r = np.sqrt(x**2+y**2)
+
+    conv_factor = (src_size/r)**(p+1)
+
+    def logplot_fp(fp: FieldPlotter, values, **kwargs) -> None:
+        fp.show_scalar_in_matplotlib(
+                np.log10(np.abs(values + 1e-15)), **kwargs)
+    if USE_MATPLOTLIB:
+        logplot_fp(fp, diff.eval(fp.points)/conv_factor, cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.colorbar()
+        plt.show()
+    1/0
+    mexp2 = t.multipole_expand(mexp, [0, 0.25])  # noqa: F841
+    lexp = t.local_expand(mexp, [3, 0])
+    lexp2 = t.local_expand(lexp, [3, 1], 3)
+
+    # diff = mexp - pt_src
+    # diff = mexp2 - pt_src
+    diff = lexp2 - pt_src
+
+    print(t.l_inf(diff, 1.2, center=lexp2.center))
+
+
+if __name__ == "__main__":
+    main()

setup.cfg

@@ -11,3 +11,4 @@ multiline-quotes = """
 [tool:pytest]
 markers =
     mpi: test distributed FMM
+    slowtest: mark a test as slow

sumpy/kernel.py

@@ -55,6 +55,7 @@
 .. autoclass:: StressletKernel
 .. autoclass:: ElasticityKernel
 .. autoclass:: LineOfCompressionKernel
+.. autoclass:: HeatKernel
 
 Derivatives
 -----------
@@ -933,6 +934,71 @@ def get_pde_as_diff_op(self):
         return laplacian(w)
 
 
+class HeatKernel(ExpressionKernel):
+    r"""The Green's function for the heat equation given by
+    :math:`e^{-r^2/{4 \alpha t^d}}/\sqrt{(4 \pi \alpha)^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.
+    """
+    init_arg_names = ("spatial_dims", "heat_alpha_name")
+
+    def __init__(self, spatial_dims, heat_alpha_name="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,
+                is_complex_valued=False)
+
+        self.heat_alpha_name = heat_alpha_name
+
+    def __getinitargs__(self):
+        return (self.dim - 1, self.heat_alpha_name)
+
+    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))
+
+    def __repr__(self):
+        return f"HeatKnl{self.dim - 1}D"
+
+    def get_args(self):
+        return [
+                KernelArgument(
+                    loopy_arg=lp.ValueArg(self.heat_alpha_name, np.float64),
+                    )]
+
+    mapper_method = "map_heat_kernel"
+
+    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)
+
+    def get_pde_as_diff_op(self):
+        from sumpy.expansion.diff_op import make_identity_diff_op, laplacian, diff
+        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)) - alpha * laplacian(w)
+
+
 # }}}
 
 
@@ -1350,6 +1416,7 @@ def map_expression_kernel(self, kernel):
     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):
         return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel))
@@ -1406,6 +1473,7 @@ def map_expression_kernel(self, kernel):
     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
 
     def map_axis_target_derivative(self, kernel):
         return 1 + self.rec(kernel.inner_kernel)