Add HeatKernel

Author: isuruf

sumpy/kernel.py

@@ -933,6 +933,61 @@ def get_pde_as_diff_op(self):
         return laplacian(w)
 
 
+class HeatKernel(ExpressionKernel):
+    init_arg_names = ("dim",)
+
+    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, 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 +1405,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 +1462,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)

test/test_kernels.py

@@ -48,6 +48,7 @@
     HelmholtzKernel,
     BiharmonicKernel,
     StokesletKernel,
+    HeatKernel,
     AxisTargetDerivative,
     DirectionalSourceDerivative)
 
@@ -278,6 +279,8 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative
             extra_kwargs["k"] = 0.2
     if isinstance(base_knl, StokesletKernel):
         extra_kwargs["mu"] = 0.2
+    if isinstance(base_knl, HeatKernel):
+        extra_kwargs["alpha"] = 1.0
 
     if with_source_derivative:
         knl = DirectionalSourceDerivative(base_knl, "dir_vec")
@@ -306,12 +309,21 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative
         h_values = [1/2, 1/3, 1/5]
 
     rng = np.random.default_rng(19)
-    center = np.array([2, 1, 0][:knl.dim], np.float64)
-    sources = actx.from_numpy(
-        0.7 * (-0.5 + rng.random((knl.dim, nsources), dtype=np.float64))
-        + center[:, np.newaxis])
+
+    if isinstance(base_knl, HeatKernel):
+        # Setup sources so that there are no negative time intervals
+        center = np.array([0, 0, 0, 0.5][-knl.dim:], np.float64)
+        sources = (-0.5 + rng.random((knl.dim, nsources), dtype=np.float64)
+            + center[:, np.newaxis])
+        loc_center = np.array([0.0, 0.0, 0.0, 6.0][-knl.dim:]) + center
+    else:
+        center = np.array([2, 1, 0][-knl.dim:], np.float64)
+        sources = 0.7 * (-0.5 + rng.random((knl.dim, nsources), dtype=np.float64)
+            + center[:, np.newaxis])
+        loc_center = np.array([5.5, 0.0, 0.0][-knl.dim:]) + center
 
     strengths = actx.from_numpy(np.ones(nsources, dtype=np.float64) / nsources)
+    sources = actx.from_numpy(sources)
 
     source_boxes = actx.from_numpy(np.array([0], dtype=np.int32))
     box_source_starts = actx.from_numpy(np.array([0], dtype=np.int32))
@@ -328,15 +340,12 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative
 
     for h in h_values:
         if issubclass(expn_class, LocalExpansionBase):
-            loc_center = np.array([5.5, 0.0, 0.0][:knl.dim]) + center
             centers = np.array(loc_center, dtype=np.float64).reshape(knl.dim, 1)
             fp = FieldPlotter(loc_center, extent=h, npoints=res)
         else:
-            eval_center = np.array([1/h, 0.0, 0.0][:knl.dim]) + center
+            eval_center = np.array([0.0, 0.0, 0.0, 1/h][-knl.dim:]) + center
             fp = FieldPlotter(eval_center, extent=0.1, npoints=res)
-            centers = (np.array([0.0, 0.0, 0.0][:knl.dim],
-                                dtype=np.float64).reshape(knl.dim, 1)
-                        + center[:, np.newaxis])
+            centers = center[:, np.newaxis]
 
         centers = actx.from_numpy(centers)
         targets = actx.from_numpy(make_obj_array(fp.points))
@@ -434,7 +443,10 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative
     if issubclass(expn_class, LocalExpansionBase):
         tgt_order_grad = tgt_order - 1
         slack = 0.7
-        grad_slack = 0.5
+        if isinstance(base_knl, HeatKernel):
+            grad_slack = 0.8
+        else:
+            grad_slack = 0.5
     else:
         tgt_order_grad = tgt_order + 1
 
@@ -445,6 +457,10 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative
             slack += 1
             grad_slack += 1
 
+        if isinstance(base_knl, HeatKernel):
+             slack += 0.5
+             grad_slack += 1.5
+
     if isinstance(knl, DirectionalSourceDerivative):
         slack += 1
         grad_slack += 2