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
 

.gitlab-ci.yml

@@ -22,6 +22,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:
@@ -38,6 +39,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:
@@ -56,6 +58,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:
@@ -72,6 +75,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-local.py

@@ -0,0 +1,85 @@
+import numpy as np
+
+import pyopencl as cl
+
+from sumpy import toys
+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 = toys.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 = 0.1
+    pt_src = toys.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:
+        toys.logplot(fp, pt_src, cmap="jet", aspect=8)
+        plt.colorbar()
+        plt.show()
+
+
+    p = 5
+    center = np.array([0, 1], dtype=np.float64)
+    mexp = toys.local_expand(pt_src, center, p)
+    diff = mexp - pt_src
+
+    dist = fp.points - center[:, None]
+    r = np.sqrt(dist[0]**2 + dist[1]**2)
+
+    error_model = 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:
+        plt.subplot(131)
+        logplot_fp(fp, error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.subplot(132)
+        logplot_fp(fp, diff.eval(fp.points), cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.subplot(133)
+        logplot_fp(fp, diff.eval(fp.points)/error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.colorbar()
+        plt.show()
+    1/0
+    mexp2 = toys.multipole_expand(mexp, [0, 0.25])  # noqa: F841
+    lexp = toys.local_expand(mexp, [3, 0])
+    lexp2 = toys.local_expand(lexp, [3, 1], 3)
+
+    # diff = mexp - pt_src
+    # diff = mexp2 - pt_src
+    diff = lexp2 - pt_src
+
+    print(toys.l_inf(diff, 1.2, center=lexp2.center))
+
+
+if __name__ == "__main__":
+    main()

examples/heat-mpole.py

@@ -0,0 +1,87 @@
+import numpy as np
+
+import pyopencl as cl
+
+from sumpy import toys
+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 = toys.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 = 0.1
+    pt_src = toys.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:
+        toys.logplot(fp, pt_src, cmap="jet", aspect=8)
+        plt.colorbar()
+        plt.show()
+
+
+    p = 5
+    mexp = toys.multipole_expand(pt_src, [0, 0], p)
+    diff = mexp - pt_src
+
+    x, t = fp.points
+
+    delta = 4 * t
+    conv_factor = src_size / np.sqrt(2*delta)
+
+    error_model = conv_factor**(p+1)*np.exp(-(x/np.sqrt(delta))**2/2)
+    #error_model = conv_factor**(p+1)/(1-conv_factor)*np.exp(-(x/np.sqrt(delta))**2)
+
+    def logplot_fp(fp: FieldPlotter, values, **kwargs) -> None:
+        fp.show_scalar_in_matplotlib(
+                np.log10(np.abs(values + 1e-15)), **kwargs)
+    if USE_MATPLOTLIB:
+        plt.subplot(131)
+        logplot_fp(fp, error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.subplot(132)
+        logplot_fp(fp, diff.eval(fp.points), cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.subplot(133)
+        logplot_fp(fp, diff.eval(fp.points)/error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
+        plt.colorbar()
+        plt.show()
+    1/0
+    mexp2 = toys.multipole_expand(mexp, [0, 0.25])  # noqa: F841
+    lexp = toys.local_expand(mexp, [3, 0])
+    lexp2 = toys.local_expand(lexp, [3, 1], 3)
+
+    # diff = mexp - pt_src
+    # diff = mexp2 - pt_src
+    diff = lexp2 - pt_src
+
+    print(toys.l_inf(diff, 1.2, center=lexp2.center))
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -149,6 +149,7 @@ extend-exclude = [
 [tool.pytest.ini_options]
 markers = [
     "mpi: tests distributed FMM",
+    "slowtest: mark a test as slow",
 ]
 
 [tool.basedpyright]

sumpy/e2e.py

@@ -324,7 +324,12 @@
                 [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)]
 

sumpy/expansion/m2l.py

@@ -885,10 +885,6 @@ def loopy_translate(self,
             vector = translation_classes_dependent_data[icoeff_src]
             expr = toeplitz_first_row * vector
 
-            domains.append(f"{{[icoeff_src]: icoeff_tgt<=icoeff_src<{ncoeff_src} }}")
-
-        expr = src_coeffs[icoeff_tgt] * translation_classes_dependent_data[icoeff_tgt]
-
         insns = [
             lp.Assignment(
                 assignee=tgt_coeffs[icoeff_tgt],

sumpy/kernel.py

@@ -106,6 +106,9 @@
 .. autoclass:: LineOfCompressionKernel
     :show-inheritance:
     :members: mapper_method
+.. autoclass:: HeatKernel
+    :show-inheritance:
+    :members: mapper_method
 
 Derivatives
 -----------
@@ -1084,6 +1087,76 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator:
         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
+
+
 # }}}
 
 
@@ -1554,6 +1627,7 @@ def map_expression_kernel(self, kernel: ExpressionKernel) -> 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: AxisTargetDerivative) -> Kernel:
         return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel))
@@ -1626,6 +1700,7 @@ def map_expression_kernel(self, kernel: ExpressionKernel) -> int:
     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: