diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 893a8b7d0..4594927ba 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 85b7eb079..8ecddbc5e 100644 --- a/.gitlab-ci.yml +++ b/.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: diff --git a/examples/heat-local.py b/examples/heat-local.py new file mode 100644 index 000000000..fca53353e --- /dev/null +++ b/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() diff --git a/examples/heat-mpole.py b/examples/heat-mpole.py new file mode 100644 index 000000000..7bc4cc6ac --- /dev/null +++ b/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() diff --git a/setup.cfg b/setup.cfg index 039e45f21..6c738d869 100644 --- a/setup.cfg +++ b/setup.cfg @@ -11,3 +11,4 @@ multiline-quotes = """ [tool:pytest] markers = mpi: test distributed FMM + slowtest: mark a test as slow diff --git a/sumpy/e2e.py b/sumpy/e2e.py index 0e218024c..e2e1218f7 100644 --- a/sumpy/e2e.py +++ b/sumpy/e2e.py @@ -332,7 +332,12 @@ def get_translation_loopy_insns(self, result_dtype): [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)] diff --git a/sumpy/expansion/m2l.py b/sumpy/expansion/m2l.py index 9a75f8ea6..43beee582 100644 --- a/sumpy/expansion/m2l.py +++ b/sumpy/expansion/m2l.py @@ -713,9 +713,6 @@ def loopy_translate(self, tgt_expansion, src_expansion): 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], diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 6c50350b3..f3af60bcc 100644 --- a/sumpy/kernel.py +++ b/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}}/\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. + """ + 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) diff --git a/sumpy/toys.py b/sumpy/toys.py index a794ff152..05e8fc43c 100644 --- a/sumpy/toys.py +++ b/sumpy/toys.py @@ -102,7 +102,7 @@ def __init__(self, cl_context: pyopencl.Context, kernel: Kernel, local_expn_class=None, expansion_factory=None, extra_source_kwargs=None, - extra_kernel_kwargs=None): + extra_kernel_kwargs=None, m2l_use_fft=None): self.cl_context = cl_context self.queue = cl.CommandQueue(self.cl_context) self.kernel = kernel @@ -116,15 +116,23 @@ def __init__(self, cl_context: pyopencl.Context, kernel: Kernel, mpole_expn_class = \ expansion_factory.get_multipole_expansion_class(kernel) if local_expn_class is None: - from sumpy.expansion.m2l import NonFFTM2LTranslationClassFactory + from sumpy.expansion.m2l import (NonFFTM2LTranslationClassFactory, + FFTM2LTranslationClassFactory) + if m2l_use_fft: + m2l_translation_class_factory = FFTM2LTranslationClassFactory() + else: + m2l_translation_class_factory = NonFFTM2LTranslationClassFactory() local_expn_class = \ expansion_factory.get_local_expansion_class(kernel) - m2l_translation_class_factory = NonFFTM2LTranslationClassFactory() m2l_translation_class = \ m2l_translation_class_factory.get_m2l_translation_class( kernel, local_expn_class) local_expn_class = partial(local_expn_class, m2l_translation=m2l_translation_class()) + elif m2l_use_fft is not None: + raise ValueError("local_expn_class and m2l_use_fft are both supplied. " + "Use only one of these arguments") + self.mpole_expn_class = mpole_expn_class self.local_expn_class = local_expn_class @@ -181,10 +189,51 @@ def get_m2m(self, from_order, to_order): self.mpole_expn_class(self.no_target_deriv_kernel, from_order), self.mpole_expn_class(self.no_target_deriv_kernel, to_order)) + @memoize_method + def use_translation_classes_dependent_data(self): + l_expn = self.local_expn_class(self.no_target_deriv_kernel, 2) + return l_expn.m2l_translation.use_preprocessing + + @memoize_method + def use_fft(self): + l_expn = self.local_expn_class(self.no_target_deriv_kernel, 2) + return l_expn.m2l_translation.use_fft + @memoize_method def get_m2l(self, from_order, to_order): - from sumpy import E2EFromCSR - return E2EFromCSR(self.cl_context, + from sumpy import E2EFromCSR, M2LUsingTranslationClassesDependentData + if self.use_translation_classes_dependent_data(): + m2l_class = M2LUsingTranslationClassesDependentData + else: + m2l_class = E2EFromCSR + return m2l_class(self.cl_context, + self.mpole_expn_class(self.no_target_deriv_kernel, from_order), + self.local_expn_class(self.no_target_deriv_kernel, to_order)) + + @memoize_method + def get_m2l_translation_class_dependent_data_kernel(self, from_order, to_order): + from sumpy import M2LGenerateTranslationClassesDependentData + return M2LGenerateTranslationClassesDependentData(self.cl_context, + self.mpole_expn_class(self.no_target_deriv_kernel, from_order), + self.local_expn_class(self.no_target_deriv_kernel, to_order)) + + @memoize_method + def get_m2l_expansion_size(self, from_order, to_order): + m_expn = self.mpole_expn_class(self.no_target_deriv_kernel, from_order) + l_expn = self.local_expn_class(self.no_target_deriv_kernel, to_order) + return l_expn.m2l_translation.preprocess_multipole_nexprs(l_expn, m_expn) + + @memoize_method + def get_m2l_preprocess_mpole_kernel(self, from_order, to_order): + from sumpy import M2LPreprocessMultipole + return M2LPreprocessMultipole(self.cl_context, + self.mpole_expn_class(self.no_target_deriv_kernel, from_order), + self.local_expn_class(self.no_target_deriv_kernel, to_order)) + + @memoize_method + def get_m2l_postprocess_local_kernel(self, from_order, to_order): + from sumpy import M2LPostprocessLocal + return M2LPostprocessLocal(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, from_order), self.local_expn_class(self.no_target_deriv_kernel, to_order)) @@ -267,7 +316,8 @@ def _e2p(psource, targets, e2p): return pot.get(queue) -def _e2e(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs): +def _e2e(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs, + extra_kernel_kwargs): toy_ctx = psource.toy_ctx queue = toy_ctx.queue @@ -290,31 +340,124 @@ def _e2e(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs): ], dtype=np.float64).T.copy() ) + coeffs = cl.array.to_device(queue, np.array([psource.coeffs])) + args = { + "queue": toy_ctx.queue, + "src_expansions": coeffs, + "src_base_ibox": 0, + "tgt_base_ibox": 0, + "ntgt_level_boxes": 2, + "target_boxes": target_boxes, + "src_box_starts": src_box_starts, + "src_box_lists": src_box_lists, + "centers": centers, + "src_rscale": psource.rscale, + "tgt_rscale": to_rscale, + **extra_kernel_kwargs, + **toy_ctx.extra_kernel_kwargs, + } + + evt, (to_coeffs,) = e2e(**args) - evt, (to_coeffs,) = e2e( - toy_ctx.queue, - src_expansions=coeffs, - src_base_ibox=0, - tgt_base_ibox=0, - ntgt_level_boxes=2, + return expn_class( + toy_ctx, to_center, to_rscale, to_order, to_coeffs[1].get(queue), + derived_from=psource, **expn_kwargs) - target_boxes=target_boxes, - src_box_starts=src_box_starts, - src_box_lists=src_box_lists, - centers=centers, +def _m2l(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs, + translation_classes_kwargs): + toy_ctx = psource.toy_ctx + queue = toy_ctx.queue - src_rscale=psource.rscale, - tgt_rscale=to_rscale, - **toy_ctx.extra_kernel_kwargs) + coeffs = cl.array.to_device(queue, np.array([psource.coeffs])) - return expn_class( - toy_ctx, to_center, to_rscale, to_order, to_coeffs[1].get(queue), + m2l_use_translation_classes_dependent_data = \ + toy_ctx.use_translation_classes_dependent_data() + + if m2l_use_translation_classes_dependent_data: + data_kernel = translation_classes_kwargs["data_kernel"] + preprocess_kernel = translation_classes_kwargs["preprocess_kernel"] + postprocess_kernel = translation_classes_kwargs["postprocess_kernel"] + expn_size = translation_classes_kwargs["m2l_expn_size"] + + # Preprocess the mpole expansion + preprocessed_src_expansions = cl.array.zeros(queue, (1, expn_size), + dtype=np.complex128) + evt, _ = preprocess_kernel(queue, + src_expansions=coeffs, + preprocessed_src_expansions = preprocessed_src_expansions, + src_rscale=np.float64(psource.rscale), + **toy_ctx.extra_kernel_kwargs) + + if toy_ctx.use_fft: + from sumpy.tools import run_opencl_fft, get_opencl_fft_app, get_native_event + fft_app = get_opencl_fft_app(queue, (expn_size,), + dtype=preprocessed_src_expansions.dtype, inverse=False) + ifft_app = get_opencl_fft_app(queue, (expn_size,), + dtype=preprocessed_src_expansions.dtype, inverse=True) + + evt, preprocessed_src_expansions = run_opencl_fft(fft_app, queue, + preprocessed_src_expansions, inverse=False, wait_for=[evt]) + + # Compute translation classes data + m2l_translation_classes_lists = cl.array.to_device(queue, np.array([0], dtype=np.int32)) + dist = np.array(to_center - psource.center, dtype=np.float64) + dim = toy_ctx.kernel.dim + m2l_translation_vectors = cl.array.to_device(queue, dist.reshape(dim, 1)) + m2l_translation_classes_dependent_data = cl.array.zeros(queue, + (1, expn_size), dtype=np.complex128) + + evt, _ = data_kernel( + queue, + src_rscale=np.float64(psource.rscale), + ntranslation_classes=1, + translation_classes_level_start=0, + m2l_translation_vectors=m2l_translation_vectors, + m2l_translation_classes_dependent_data=( + m2l_translation_classes_dependent_data), + ntranslation_vectors=1, + wait_for=[get_native_event(evt)], + **toy_ctx.extra_kernel_kwargs) + + if toy_ctx.use_fft: + evt, m2l_translation_classes_dependent_data = run_opencl_fft(fft_app, queue, + m2l_translation_classes_dependent_data, inverse=False, wait_for=[evt]) + + ret = _e2e(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs, + {"src_expansions": preprocessed_src_expansions, + "m2l_translation_classes_lists": m2l_translation_classes_lists, + "m2l_translation_classes_dependent_data": + m2l_translation_classes_dependent_data, + "translation_classes_level_start": 0, + }) + + # Postprocess the local expansion + local_before = cl.array.to_device(queue, np.array([ret.coeffs])) + to_coeffs = cl.array.zeros(queue, (1, len(data_kernel.tgt_expansion)), + dtype=coeffs.dtype) + + if toy_ctx.use_fft: + evt, local_before = run_opencl_fft(ifft_app, queue, + local_before, inverse=True, wait_for=[get_native_event(evt)]) + + evt, _ = postprocess_kernel(queue=queue, + tgt_expansions_before_postprocessing=local_before, + tgt_expansions=to_coeffs, + src_rscale=np.float64(psource.rscale), + tgt_rscale=np.float64(to_rscale), + wait_for=[get_native_event(evt)], + **toy_ctx.extra_kernel_kwargs) + + return expn_class( + toy_ctx, to_center, to_rscale, to_order, to_coeffs.get(queue)[0], derived_from=psource, **expn_kwargs) + else: + ret = _e2e(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs, {}) -# }}} + return ret +# }}} # {{{ potential source classes @@ -596,7 +739,7 @@ def multipole_expand(psource: PotentialSource, return _e2e(psource, center, rscale, order, psource.toy_ctx.get_m2m(psource.order, order), - MultipoleExpansion, expn_kwargs) + MultipoleExpansion, expn_kwargs, {}) else: raise TypeError(f"do not know how to expand '{type(psource).__name__}'") @@ -617,9 +760,28 @@ def local_expand( if order is None: order = psource.order - return _e2e(psource, center, rscale, order, - psource.toy_ctx.get_m2l(psource.order, order), - LocalExpansion, expn_kwargs) + toy_ctx = psource.toy_ctx + translation_classes_kwargs = {} + m2l_use_translation_classes_dependent_data = \ + toy_ctx.use_translation_classes_dependent_data() + + if m2l_use_translation_classes_dependent_data: + data_kernel = toy_ctx.get_m2l_translation_class_dependent_data_kernel( + psource.order, order) + preprocess_kernel = toy_ctx.get_m2l_preprocess_mpole_kernel( + psource.order, order) + postprocess_kernel = toy_ctx.get_m2l_postprocess_local_kernel( + psource.order, order) + translation_classes_kwargs["data_kernel"] = data_kernel + translation_classes_kwargs["preprocess_kernel"] = preprocess_kernel + translation_classes_kwargs["postprocess_kernel"] = postprocess_kernel + translation_classes_kwargs["m2l_expn_size"] = \ + toy_ctx.get_m2l_expansion_size(psource.order, order) + + return _m2l(psource, center, rscale, order, + toy_ctx.get_m2l(psource.order, order), + LocalExpansion, expn_kwargs, + translation_classes_kwargs) elif isinstance(psource, LocalExpansion): if order is None: @@ -627,7 +789,7 @@ def local_expand( return _e2e(psource, center, rscale, order, psource.toy_ctx.get_l2l(psource.order, order), - LocalExpansion, expn_kwargs) + LocalExpansion, expn_kwargs, {}) else: raise TypeError(f"do not know how to expand '{type(psource).__name__}'") diff --git a/test/test_kernels.py b/test/test_kernels.py index 2ddbd9c5f..791cf656d 100644 --- a/test/test_kernels.py +++ b/test/test_kernels.py @@ -48,6 +48,7 @@ HelmholtzKernel, BiharmonicKernel, StokesletKernel, + HeatKernel, AxisTargetDerivative, DirectionalSourceDerivative) @@ -239,6 +240,15 @@ def test_p2e_multiple(actx_factory, base_knl, expn_class): (LaplaceKernel(2), LinearPDEConformingVolumeTaylorLocalExpansion), (LaplaceKernel(2), LinearPDEConformingVolumeTaylorMultipoleExpansion), + (HeatKernel(1), LinearPDEConformingVolumeTaylorLocalExpansion), + (HeatKernel(1), LinearPDEConformingVolumeTaylorMultipoleExpansion), + (HeatKernel(2), LinearPDEConformingVolumeTaylorLocalExpansion), + (HeatKernel(2), LinearPDEConformingVolumeTaylorMultipoleExpansion), + pytest.param(HeatKernel(3), LinearPDEConformingVolumeTaylorLocalExpansion, + marks=pytest.mark.slowtest), + pytest.param(HeatKernel(3), LinearPDEConformingVolumeTaylorMultipoleExpansion, + marks=pytest.mark.slowtest), + (HelmholtzKernel(2), VolumeTaylorMultipoleExpansion), (HelmholtzKernel(2), VolumeTaylorLocalExpansion), (HelmholtzKernel(2), LinearPDEConformingVolumeTaylorLocalExpansion), @@ -278,6 +288,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 +318,23 @@ 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 + varying_axis = -1 + 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 + varying_axis = 0 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)) @@ -321,22 +344,21 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative extra_source_kwargs = extra_kwargs.copy() if isinstance(knl, DirectionalSourceDerivative): alpha = np.linspace(0, 2*np.pi, nsources, np.float64) - dir_vec = np.vstack([np.cos(alpha), np.sin(alpha)]) + dir_vec = np.vstack( + [np.cos(alpha), np.sin(alpha), alpha, alpha**2][:knl.dim]) extra_source_kwargs["dir_vec"] = actx.from_numpy(dir_vec) from sumpy.visualization import FieldPlotter 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 = center.copy() + eval_center[varying_axis] = 1/h 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 +456,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 +470,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 += 1.0 + grad_slack += 2.5 + if isinstance(knl, DirectionalSourceDerivative): slack += 1 grad_slack += 2 @@ -471,6 +500,12 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative # {{{ test_translations @pytest.mark.parametrize("knl, local_expn_class, mpole_expn_class", [ + (HeatKernel(1), LinearPDEConformingVolumeTaylorLocalExpansion, + LinearPDEConformingVolumeTaylorMultipoleExpansion), + (HeatKernel(2), LinearPDEConformingVolumeTaylorLocalExpansion, + LinearPDEConformingVolumeTaylorMultipoleExpansion), + (HeatKernel(3), LinearPDEConformingVolumeTaylorLocalExpansion, + LinearPDEConformingVolumeTaylorMultipoleExpansion), (LaplaceKernel(2), VolumeTaylorLocalExpansion, VolumeTaylorMultipoleExpansion), (LaplaceKernel(2), LinearPDEConformingVolumeTaylorLocalExpansion, LinearPDEConformingVolumeTaylorMultipoleExpansion), @@ -506,10 +541,12 @@ def test_translations(actx_factory, knl, local_expn_class, mpole_expn_class, extra_kwargs["k"] = 0.05 if isinstance(knl, StokesletKernel): extra_kwargs["mu"] = 0.05 + if isinstance(knl, HeatKernel): + extra_kwargs["alpha"] = 0.1 # Just to make sure things also work away from the origin rng = np.random.default_rng(18) - origin = np.array([2, 1, 0][:knl.dim], np.float64) + origin = np.array([0, 0, 1, 2][-knl.dim:], np.float64) sources = actx.from_numpy( 0.7 * (-0.5 + rng.random((knl.dim, nsources), dtype=np.float64)) + origin[:, np.newaxis]) @@ -523,19 +560,19 @@ def test_translations(actx_factory, knl, local_expn_class, mpole_expn_class, from sumpy.visualization import FieldPlotter - eval_offset = np.array([5.5, 0.0, 0][:knl.dim]) + eval_offset = np.array([0.0, 0.0, 0.0, 5.5][-knl.dim:]) fp = FieldPlotter(eval_offset + origin, extent=0.3, npoints=res) centers = actx.from_numpy((np.array( [ # box 0: particles, first mpole here - [0, 0, 0][:knl.dim], + [0, 0, 0, 0][-knl.dim:], # box 1: second mpole here - np.array([-0.2, 0.1, 0][:knl.dim], np.float64), + np.array([0, 0, 0.1, -0.2][-knl.dim:], np.float64), # box 2: first local here - eval_offset + np.array([0.3, -0.2, 0][:knl.dim], np.float64), + eval_offset + np.array([0, 0, -0.2, 0.3][-knl.dim:], np.float64), # box 3: second local and eval here eval_offset @@ -544,7 +581,7 @@ def test_translations(actx_factory, knl, local_expn_class, mpole_expn_class, del eval_offset - if knl.dim == 2: + if knl.dim == 2 and not isinstance(knl, HeatKernel): orders = [2, 3, 4] else: orders = [3, 4, 5] diff --git a/test/test_misc.py b/test/test_misc.py index 9b7257812..5168aa6b3 100644 --- a/test/test_misc.py +++ b/test/test_misc.py @@ -44,6 +44,7 @@ StressletKernel, ElasticityKernel, LineOfCompressionKernel, + HeatKernel, ExpressionKernel) from sumpy.expansion.diff_op import ( make_identity_diff_op, concat, as_scalar_pde, diff, @@ -110,6 +111,9 @@ def nderivs(self): KernelInfo(ElasticityKernel(3, 0, 0), mu=5, nu=0.2), KernelInfo(LineOfCompressionKernel(3, 0), mu=5, nu=0.2), KernelInfo(LineOfCompressionKernel(3, 1), mu=5, nu=0.2), + KernelInfo(HeatKernel(1), alpha=0.1), + KernelInfo(HeatKernel(2), alpha=0.1), + KernelInfo(HeatKernel(3), alpha=0.1), ]) def test_pde_check_kernels(actx_factory, knl_info, order=5): actx = actx_factory() @@ -129,7 +133,7 @@ def test_pde_check_kernels(actx_factory, knl_info, order=5): eoc_rec = EOCRecorder() for h in [0.1, 0.05, 0.025]: - cp = CalculusPatch(np.array([1, 0, 0])[:dim], h=h, order=order) + cp = CalculusPatch(np.array([0, 0, 0, 1])[-dim:], h=h, order=order) pot = pt_src.eval(cp.points) pde = knl_info.pde_func(cp, pot) @@ -230,6 +234,7 @@ class P2E2E2PTestCase: expansion1: Callable[..., Any] expansion2: Callable[..., Any] conv_factor: str + m2l_use_fft: bool = False @property def dim(self): @@ -266,6 +271,17 @@ def dim(self): expansion1=t.multipole_expand, expansion2=t.local_expand, conv_factor="norm(t-c2)/(norm(c2-c1)-norm(c1-s))"), + + # multipole to local, 3D with FFT + P2E2E2PTestCase( + source=np.array([-2., 2., 1.]), + center1=np.array([-2., 5., 3.]), + center2=np.array([0., 0., 0.]), + target=np.array([0., 0., -1]), + expansion1=t.multipole_expand, + expansion2=t.local_expand, + m2l_use_fft=True, + conv_factor="norm(t-c2)/(norm(c2-c1)-norm(c1-s))"), ) # }}} @@ -297,12 +313,14 @@ def test_toy_p2e2e2p(actx_factory, case): raise ValueError( f"convergence factor not in valid range: {case_conv_factor}") - from sumpy.expansion import VolumeTaylorExpansionFactory + from sumpy.expansion import DefaultExpansionFactory actx = actx_factory() ctx = t.ToyContext(actx.context, - LaplaceKernel(dim), - expansion_factory=VolumeTaylorExpansionFactory()) + HeatKernel(dim - 1), + expansion_factory=DefaultExpansionFactory(), + m2l_use_fft=case.m2l_use_fft, + extra_kernel_kwargs={"alpha": 0.1}) errors = []