Skip to content

Commit e452098

Browse files
isurufinducer
authored andcommitted
Add heat kernel
Squashed from #185
1 parent 69eacc9 commit e452098

11 files changed

Lines changed: 347 additions & 31 deletions

File tree

.github/workflows/ci.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ jobs:
6868
run: |
6969
grep -v symengine .test-conda-env-py3.yml > .test-conda-env.yml
7070
CONDA_ENVIRONMENT=.test-conda-env.yml
71+
export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
7172
curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
7273
. ./build-and-test-py-project-within-miniconda.sh
7374
@@ -89,6 +90,7 @@ jobs:
8990
- uses: actions/checkout@v6
9091
- name: "Main Script"
9192
run: |
93+
export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
9294
curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
9395
. ./build-and-test-py-project-within-miniconda.sh
9496

.gitlab-ci.yml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ Pytest POCL:
2222
- export PY_EXE=python3
2323
- export PYOPENCL_TEST=portable:pthread
2424
- export EXTRA_INSTALL="pybind11 numpy mako mpi4py"
25+
- export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
2526
- curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
2627
- ". ./build-and-test-py-project.sh"
2728
tags:
@@ -38,6 +39,7 @@ Pytest Titan V:
3839
- py_version=3
3940
- export PYOPENCL_TEST=nvi:titan
4041
- EXTRA_INSTALL="pybind11 numpy mako mpi4py"
42+
- export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
4143
- curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
4244
- ". ./build-and-test-py-project.sh"
4345
tags:
@@ -56,6 +58,7 @@ Pytest Conda:
5658
- export SUMPY_NO_CACHE=1
5759
- export SUMPY_FORCE_SYMBOLIC_BACKEND=symengine
5860
- export PYOPENCL_TEST=portable:pthread
61+
- export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
5962
- curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
6063
- ". ./build-and-test-py-project-within-miniconda.sh"
6164
tags:
@@ -72,6 +75,7 @@ Pytest POCL Titan V:
7275
- export SUMPY_NO_CACHE=1
7376
- export SUMPY_FORCE_SYMBOLIC_BACKEND=symengine
7477
- export PYOPENCL_TEST=portable:titan
78+
- export PYTEST_ADDOPTS=${PYTEST_ADDOPTS:-"-m 'not slowtest'"}
7579
- curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
7680
- ". ./build-and-test-py-project-within-miniconda.sh"
7781
tags:

examples/heat-local.py

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import numpy as np
2+
3+
import pyopencl as cl
4+
5+
from sumpy import toys
6+
from sumpy.visualization import FieldPlotter
7+
from sumpy.kernel import ( # noqa: F401
8+
YukawaKernel,
9+
HeatKernel,
10+
HelmholtzKernel,
11+
LaplaceKernel)
12+
13+
try:
14+
import matplotlib.pyplot as plt
15+
USE_MATPLOTLIB = True
16+
except ImportError:
17+
USE_MATPLOTLIB = False
18+
19+
20+
def main():
21+
from sumpy.array_context import PyOpenCLArrayContext
22+
ctx = cl.create_some_context()
23+
queue = cl.CommandQueue(ctx)
24+
actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
25+
26+
tctx = toys.ToyContext(
27+
actx.context,
28+
# LaplaceKernel(2),
29+
#YukawaKernel(2), extra_kernel_kwargs={"lam": 5},
30+
# HelmholtzKernel(2), extra_kernel_kwargs={"k": 0.3},
31+
HeatKernel(1), extra_kernel_kwargs={"alpha": 1},
32+
)
33+
34+
src_size = 0.1
35+
pt_src = toys.PointSources(
36+
tctx,
37+
np.array([
38+
src_size*(np.random.rand(50) - 0.5),
39+
np.zeros(50)]),
40+
np.random.randn(50))
41+
42+
fp = FieldPlotter([0, 0.5], extent=np.array([8, 1]))
43+
44+
if 0 and USE_MATPLOTLIB:
45+
toys.logplot(fp, pt_src, cmap="jet", aspect=8)
46+
plt.colorbar()
47+
plt.show()
48+
49+
50+
p = 5
51+
center = np.array([0, 1], dtype=np.float64)
52+
mexp = toys.local_expand(pt_src, center, p)
53+
diff = mexp - pt_src
54+
55+
dist = fp.points - center[:, None]
56+
r = np.sqrt(dist[0]**2 + dist[1]**2)
57+
58+
error_model = r**(p+1)
59+
60+
def logplot_fp(fp: FieldPlotter, values, **kwargs) -> None:
61+
fp.show_scalar_in_matplotlib(
62+
np.log10(np.abs(values + 1e-15)), **kwargs)
63+
if USE_MATPLOTLIB:
64+
plt.subplot(131)
65+
logplot_fp(fp, error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
66+
plt.subplot(132)
67+
logplot_fp(fp, diff.eval(fp.points), cmap="jet", vmin=-5, vmax=0, aspect=8)
68+
plt.subplot(133)
69+
logplot_fp(fp, diff.eval(fp.points)/error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
70+
plt.colorbar()
71+
plt.show()
72+
1/0
73+
mexp2 = toys.multipole_expand(mexp, [0, 0.25]) # noqa: F841
74+
lexp = toys.local_expand(mexp, [3, 0])
75+
lexp2 = toys.local_expand(lexp, [3, 1], 3)
76+
77+
# diff = mexp - pt_src
78+
# diff = mexp2 - pt_src
79+
diff = lexp2 - pt_src
80+
81+
print(toys.l_inf(diff, 1.2, center=lexp2.center))
82+
83+
84+
if __name__ == "__main__":
85+
main()

examples/heat-mpole.py

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
import numpy as np
2+
3+
import pyopencl as cl
4+
5+
from sumpy import toys
6+
from sumpy.visualization import FieldPlotter
7+
from sumpy.kernel import ( # noqa: F401
8+
YukawaKernel,
9+
HeatKernel,
10+
HelmholtzKernel,
11+
LaplaceKernel)
12+
13+
try:
14+
import matplotlib.pyplot as plt
15+
USE_MATPLOTLIB = True
16+
except ImportError:
17+
USE_MATPLOTLIB = False
18+
19+
20+
def main():
21+
from sumpy.array_context import PyOpenCLArrayContext
22+
ctx = cl.create_some_context()
23+
queue = cl.CommandQueue(ctx)
24+
actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
25+
26+
tctx = toys.ToyContext(
27+
actx.context,
28+
# LaplaceKernel(2),
29+
#YukawaKernel(2), extra_kernel_kwargs={"lam": 5},
30+
# HelmholtzKernel(2), extra_kernel_kwargs={"k": 0.3},
31+
HeatKernel(1), extra_kernel_kwargs={"alpha": 1},
32+
)
33+
34+
src_size = 0.1
35+
pt_src = toys.PointSources(
36+
tctx,
37+
np.array([
38+
src_size*(np.random.rand(50) - 0.5),
39+
np.zeros(50)]),
40+
np.random.randn(50))
41+
42+
fp = FieldPlotter([0, 0.5], extent=np.array([8, 1]))
43+
44+
if 0 and USE_MATPLOTLIB:
45+
toys.logplot(fp, pt_src, cmap="jet", aspect=8)
46+
plt.colorbar()
47+
plt.show()
48+
49+
50+
p = 5
51+
mexp = toys.multipole_expand(pt_src, [0, 0], p)
52+
diff = mexp - pt_src
53+
54+
x, t = fp.points
55+
56+
delta = 4 * t
57+
conv_factor = src_size / np.sqrt(2*delta)
58+
59+
error_model = conv_factor**(p+1)*np.exp(-(x/np.sqrt(delta))**2/2)
60+
#error_model = conv_factor**(p+1)/(1-conv_factor)*np.exp(-(x/np.sqrt(delta))**2)
61+
62+
def logplot_fp(fp: FieldPlotter, values, **kwargs) -> None:
63+
fp.show_scalar_in_matplotlib(
64+
np.log10(np.abs(values + 1e-15)), **kwargs)
65+
if USE_MATPLOTLIB:
66+
plt.subplot(131)
67+
logplot_fp(fp, error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
68+
plt.subplot(132)
69+
logplot_fp(fp, diff.eval(fp.points), cmap="jet", vmin=-5, vmax=0, aspect=8)
70+
plt.subplot(133)
71+
logplot_fp(fp, diff.eval(fp.points)/error_model, cmap="jet", vmin=-5, vmax=0, aspect=8)
72+
plt.colorbar()
73+
plt.show()
74+
1/0
75+
mexp2 = toys.multipole_expand(mexp, [0, 0.25]) # noqa: F841
76+
lexp = toys.local_expand(mexp, [3, 0])
77+
lexp2 = toys.local_expand(lexp, [3, 1], 3)
78+
79+
# diff = mexp - pt_src
80+
# diff = mexp2 - pt_src
81+
diff = lexp2 - pt_src
82+
83+
print(toys.l_inf(diff, 1.2, center=lexp2.center))
84+
85+
86+
if __name__ == "__main__":
87+
main()

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ extend-exclude = [
154154
[tool.pytest.ini_options]
155155
markers = [
156156
"mpi: tests distributed FMM",
157+
"slowtest: mark a test as slow",
157158
]
158159

159160
[tool.basedpyright]

sumpy/e2e.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,7 +320,12 @@ def get_translation_loopy_insns(self, result_dtype):
320320
[sym.Symbol(f"data{i}")
321321
for i in range(m2l_translation_classes_dependent_ndata)]
322322

323-
ncoeff_src = len(self.src_expansion)
323+
m2l_translation = self.tgt_expansion.m2l_translation
324+
if m2l_translation.use_preprocessing:
325+
ncoeff_src = m2l_translation.preprocess_multipole_nexprs(
326+
self.tgt_expansion, self.src_expansion)
327+
else:
328+
ncoeff_src = len(self.src_expansion)
324329

325330
src_coeff_exprs = [sym.Symbol(f"src_coeffs{i}") for i in range(ncoeff_src)]
326331

sumpy/expansion/m2l.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -879,10 +879,6 @@ def loopy_translate(self,
879879
vector = translation_classes_dependent_data[icoeff_src]
880880
expr = toeplitz_first_row * vector
881881

882-
domains.append(f"{{[icoeff_src]: icoeff_tgt<=icoeff_src<{ncoeff_src} }}")
883-
884-
expr = src_coeffs[icoeff_tgt] * translation_classes_dependent_data[icoeff_tgt]
885-
886882
insns = [
887883
lp.Assignment(
888884
assignee=tgt_coeffs[icoeff_tgt],

sumpy/kernel.py

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,9 @@
106106
.. autoclass:: LineOfCompressionKernel
107107
:show-inheritance:
108108
:members: mapper_method
109+
.. autoclass:: HeatKernel
110+
:show-inheritance:
111+
:members: mapper_method
109112
110113
Derivatives
111114
-----------
@@ -1119,6 +1122,76 @@ def get_pde_as_diff_op(self) -> LinearPDESystemOperator:
11191122
return laplacian(w)
11201123

11211124

1125+
class HeatKernel(ExpressionKernel):
1126+
r"""The Green's function for the heat equation given by
1127+
:math:`e^{-r^2/{4 \alpha t}}/\sqrt{(4 \pi \alpha t)^d}`
1128+
where :math:`d` is the number of spatial dimensions.
1129+
1130+
.. note::
1131+
1132+
This kernel cannot be used in an FMM yet and can only
1133+
be used in expansions and evaluations that occur forward
1134+
in the time dimension.
1135+
"""
1136+
heat_alpha_name: str
1137+
1138+
mapper_method: ClassVar[str] = "map_heat_kernel"
1139+
1140+
def __init__(self, spatial_dims: int, heat_alpha_name: str = "alpha"):
1141+
dim = spatial_dims + 1
1142+
d = make_sym_vector("d", dim)
1143+
t = d[-1]
1144+
r = pymbolic_real_norm_2(d[:-1])
1145+
alpha = SpatialConstant(heat_alpha_name)
1146+
expr = var("exp")(-r**2/(4 * alpha * t)) / var("sqrt")(t**(dim - 1))
1147+
scaling = 1/var("sqrt")((4*var("pi")*alpha)**(dim - 1))
1148+
1149+
super().__init__(
1150+
dim,
1151+
expression=expr,
1152+
global_scaling_const=scaling,
1153+
)
1154+
1155+
self.heat_alpha_name = heat_alpha_name
1156+
1157+
@property
1158+
@override
1159+
def is_complex_valued(self) -> bool:
1160+
return False
1161+
1162+
def update_persistent_hash(self, key_hash, key_builder):
1163+
key_hash.update(type(self).__name__.encode("utf8"))
1164+
key_builder.rec(key_hash, (self.dim, self.heat_alpha_name))
1165+
1166+
@override
1167+
def __repr__(self):
1168+
return f"HeatKnl{self.dim - 1}D"
1169+
1170+
@override
1171+
def get_args(self):
1172+
return [
1173+
KernelArgument(
1174+
loopy_arg=lp.ValueArg(self.heat_alpha_name, np.float64),
1175+
)]
1176+
1177+
def get_derivative_taker(self, dvec, rscale, sac):
1178+
"""Return a :class:`sumpy.derivative_taker.ExprDerivativeTaker` instance
1179+
that supports taking derivatives of the base kernel with respect to dvec.
1180+
"""
1181+
from sumpy.derivative_taker import ExprDerivativeTaker
1182+
return ExprDerivativeTaker(self.get_expression(dvec), dvec, rscale,
1183+
sac)
1184+
1185+
@override
1186+
def get_pde_as_diff_op(self):
1187+
from sumpy.expansion.diff_op import diff, laplacian, make_identity_diff_op
1188+
alpha = sym.Symbol(self.heat_alpha_name)
1189+
w = make_identity_diff_op(self.dim - 1, time_dependent=True)
1190+
time_diff = [0]*self.dim
1191+
time_diff[-1] = 1
1192+
return diff(w, tuple(time_diff)) - laplacian(w) * alpha
1193+
1194+
11221195
# }}}
11231196

11241197

@@ -1590,6 +1663,7 @@ def map_expression_kernel(self, kernel: ExpressionKernel) -> Kernel:
15901663
map_elasticity_kernel = map_expression_kernel
15911664
map_line_of_compression_kernel = map_expression_kernel
15921665
map_stresslet_kernel = map_expression_kernel
1666+
map_heat_kernel = map_expression_kernel
15931667

15941668
def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> Kernel:
15951669
return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel))
@@ -1662,6 +1736,7 @@ def map_expression_kernel(self, kernel: ExpressionKernel) -> int:
16621736
map_yukawa_kernel = map_expression_kernel
16631737
map_line_of_compression_kernel = map_expression_kernel
16641738
map_stresslet_kernel = map_expression_kernel
1739+
map_heat_kernel = map_expression_kernel
16651740

16661741
@override
16671742
def map_axis_target_derivative(self, kernel: AxisTargetDerivative) -> int:

0 commit comments

Comments
 (0)