diff --git a/requirements.txt b/requirements.txt index 72c6f73bc4acedb82d1a49800cdfc2d2863757c8..8204d06123ea7735a026cff46285bebb34e68af7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ numpy sympy==1.1.1 pyrsistent -git+https://github.com/inducer/pytools#egg=pytools -git+https://github.com/inducer/pymbolic#egg=pymbolic -git+https://github.com/inducer/islpy#egg=islpy -git+https://github.com/inducer/pyopencl#egg=pyopencl -git+https://github.com/inducer/boxtree#egg=boxtree -git+https://github.com/inducer/loopy#egg=loopy -git+https://github.com/inducer/pyfmmlib#egg=pyfmmlib +git+https://github.com/inducer/pytools.git#egg=pytools +git+https://github.com/inducer/pymbolic.git#egg=pymbolic +git+https://github.com/inducer/islpy.git#egg=islpy +git+https://github.com/inducer/pyopencl.git#egg=pyopencl +git+https://github.com/inducer/boxtree.git#egg=boxtree +git+https://github.com/inducer/loopy.git#egg=loopy +git+https://github.com/inducer/pyfmmlib.git#egg=pyfmmlib diff --git a/sumpy/e2e.py b/sumpy/e2e.py index fdc937ee382eaa2c8fef0da90cb278b6a691401b..52dd453eaf28192e498e0bed56217a238c3c6ecc 100644 --- a/sumpy/e2e.py +++ b/sumpy/e2e.py @@ -48,7 +48,7 @@ Expansion-to-expansion class E2EBase(KernelCacheWrapper): def __init__(self, ctx, src_expansion, tgt_expansion, - options=[], name=None, device=None): + name=None, device=None): """ :arg expansion: a subclass of :class:`sympy.expansion.ExpansionBase` :arg strength_usage: A list of integers indicating which expression @@ -79,7 +79,6 @@ class E2EBase(KernelCacheWrapper): self.ctx = ctx self.src_expansion = src_expansion self.tgt_expansion = tgt_expansion - self.options = options self.name = name or self.default_name self.device = device diff --git a/sumpy/e2p.py b/sumpy/e2p.py index 964dccb1fb37ef9d29203b4902f88bb87bd0523b..2512aca8574cac4c1815a1202c729aab1e25786f 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -44,7 +44,7 @@ Expansion-to-particle class E2PBase(KernelCacheWrapper): def __init__(self, ctx, expansion, kernels, - options=[], name=None, device=None): + name=None, device=None): """ :arg expansion: a subclass of :class:`sympy.expansion.ExpansionBase` :arg strength_usage: A list of integers indicating which expression @@ -68,7 +68,6 @@ class E2PBase(KernelCacheWrapper): self.ctx = ctx self.expansion = expansion self.kernels = kernels - self.options = options self.name = name or self.default_name self.device = device diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py index 83dc6b30ddb9c3237c30ede4059455992eb7b271..436707e1fd8f7b29900da0bbcfe6cc0f8c672c14 100644 --- a/sumpy/expansion/local.py +++ b/sumpy/expansion/local.py @@ -55,7 +55,7 @@ class LineTaylorLocalExpansion(LocalExpansionBase): def get_coefficient_identifiers(self): return list(range(self.order+1)) - def coefficients_from_source(self, avec, bvec, rscale, sac=None): + def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None): # no point in heeding rscale here--just ignore it if bvec is None: raise RuntimeError("cannot use line-Taylor expansions in a setting " @@ -66,7 +66,7 @@ class LineTaylorLocalExpansion(LocalExpansionBase): avec_line = avec + tau*bvec - line_kernel = self.kernel.get_expression(avec_line) + line_kernel = kernel.get_expression(avec_line) from sumpy.symbolic import USE_SYMENGINE @@ -75,8 +75,8 @@ class LineTaylorLocalExpansion(LocalExpansionBase): deriv_taker = MiDerivativeTaker(line_kernel, (tau,)) return [my_syntactic_subs( - self.kernel.postprocess_at_target( - self.kernel.postprocess_at_source( + kernel.postprocess_at_target( + kernel.postprocess_at_source( deriv_taker.diff(i), avec), bvec), {tau: 0}) @@ -89,8 +89,8 @@ class LineTaylorLocalExpansion(LocalExpansionBase): # # See also https://gitlab.tiker.net/inducer/pytential/merge_requests/12 - return [self.kernel.postprocess_at_target( - self.kernel.postprocess_at_source( + return [kernel.postprocess_at_target( + kernel.postprocess_at_source( line_kernel.diff("tau", i), avec), bvec) .subs("tau", 0) @@ -113,10 +113,9 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): Coefficients represent derivative values of the kernel. """ - def coefficients_from_source(self, avec, bvec, rscale, sac=None): + def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None): from sumpy.tools import MiDerivativeTaker - ppkernel = self.kernel.postprocess_at_source( - self.kernel.get_expression(avec), avec) + ppkernel = kernel.postprocess_at_source(kernel.get_expression(avec), avec) taker = MiDerivativeTaker(ppkernel, avec) return [ @@ -372,7 +371,7 @@ class _FourierBesselLocalExpansion(LocalExpansionBase): def get_coefficient_identifiers(self): return list(range(-self.order, self.order+1)) - def coefficients_from_source(self, avec, bvec, rscale, sac=None): + def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None): if not self.use_rscale: rscale = 1 @@ -384,7 +383,7 @@ class _FourierBesselLocalExpansion(LocalExpansionBase): # The coordinates are negated since avec points from source to center. source_angle_rel_center = sym.atan2(-avec[1], -avec[0]) avec_len = sym_real_norm_2(avec) - return [self.kernel.postprocess_at_source( + return [kernel.postprocess_at_source( hankel_1(c, arg_scale * avec_len) * rscale ** abs(c) * sym.exp(sym.I * c * source_angle_rel_center), avec) diff --git a/sumpy/expansion/multipole.py b/sumpy/expansion/multipole.py index 313b49bdffce8a3027e64904dee4426545c4b013..e316bdfc4cefabaf841a61d0acd6baec3be6d54b 100644 --- a/sumpy/expansion/multipole.py +++ b/sumpy/expansion/multipole.py @@ -53,9 +53,10 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): Coefficients represent the terms in front of the kernel derivatives. """ - def coefficients_from_source(self, avec, bvec, rscale, sac=None): + def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None): from sumpy.kernel import DirectionalSourceDerivative - kernel = self.kernel + if kernel is None: + kernel = self.kernel from sumpy.tools import mi_power, mi_factorial @@ -284,10 +285,13 @@ class _HankelBased2DMultipoleExpansion(MultipoleExpansionBase): def get_coefficient_identifiers(self): return list(range(-self.order, self.order+1)) - def coefficients_from_source(self, avec, bvec, rscale, sac=None): + def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None): if not self.use_rscale: rscale = 1 + if kernel is None: + kernel = self.kernel + from sumpy.symbolic import sym_real_norm_2 bessel_j = sym.Function("bessel_j") avec_len = sym_real_norm_2(avec) @@ -297,7 +301,7 @@ class _HankelBased2DMultipoleExpansion(MultipoleExpansionBase): # The coordinates are negated since avec points from source to center. source_angle_rel_center = sym.atan2(-avec[1], -avec[0]) return [ - self.kernel.postprocess_at_source( + kernel.postprocess_at_source( bessel_j(c, arg_scale * avec_len) / rscale ** abs(c) * sym.exp(sym.I * c * -source_angle_rel_center), diff --git a/sumpy/fmm.py b/sumpy/fmm.py index 39ed7f2093d44a32c24a5265cfb0ba844de2e98d..feebbf3d91d577052b6d95b44e6f7ad72c21efe8 100644 --- a/sumpy/fmm.py +++ b/sumpy/fmm.py @@ -351,7 +351,7 @@ class SumpyExpansionWrangler: def form_multipoles(self, level_start_source_box_nrs, source_boxes, - src_weights): + src_weight_vecs): mpoles = self.multipole_expansion_zeros() kwargs = self.extra_kwargs.copy() @@ -372,7 +372,7 @@ class SumpyExpansionWrangler: self.queue, source_boxes=source_boxes[start:stop], centers=self.tree.box_centers, - strengths=src_weights, + strengths=src_weight_vecs, tgt_expansions=mpoles_view, tgt_base_ibox=level_start_ibox, @@ -443,7 +443,7 @@ class SumpyExpansionWrangler: return (mpoles, SumpyTimingFuture(self.queue, events)) def eval_direct(self, target_boxes, source_box_starts, - source_box_lists, src_weights): + source_box_lists, src_weight_vecs): pot = self.output_zeros() kwargs = self.extra_kwargs.copy() @@ -457,7 +457,7 @@ class SumpyExpansionWrangler: target_boxes=target_boxes, source_box_starts=source_box_starts, source_box_lists=source_box_lists, - strength=(src_weights,), + strength=src_weight_vecs, result=pot, **kwargs) @@ -563,7 +563,7 @@ class SumpyExpansionWrangler: def form_locals(self, level_start_target_or_target_parent_box_nrs, - target_or_target_parent_boxes, starts, lists, src_weights): + target_or_target_parent_boxes, starts, lists, src_weight_vecs): local_exps = self.local_expansion_zeros() kwargs = self.extra_kwargs.copy() @@ -588,7 +588,7 @@ class SumpyExpansionWrangler: source_box_starts=starts[start:stop+1], source_box_lists=lists, centers=self.tree.box_centers, - strengths=src_weights, + strengths=src_weight_vecs, tgt_expansions=target_local_exps_view, tgt_base_ibox=target_level_start_ibox, diff --git a/sumpy/p2e.py b/sumpy/p2e.py index 58c38321d0270382ad58111ccf67421c1abc4ba8..93692e6ec47abd85a24f35a591bec0b95ab287d5 100644 --- a/sumpy/p2e.py +++ b/sumpy/p2e.py @@ -23,8 +23,9 @@ THE SOFTWARE. import numpy as np import loopy as lp from loopy.version import MOST_RECENT_LANGUAGE_VERSION +import pymbolic -from sumpy.tools import KernelCacheWrapper +from sumpy.tools import KernelCacheWrapper, KernelComputation import logging logger = logging.getLogger(__name__) @@ -44,9 +45,11 @@ Particle-to-expansion # {{{ P2E base class -class P2EBase(KernelCacheWrapper): - def __init__(self, ctx, expansion, - options=[], name=None, device=None): +class P2EBase(KernelComputation, KernelCacheWrapper): + """Common input processing for kernel computations.""" + + def __init__(self, ctx, expansion, kernels=None, + name=None, device=None, strength_usage=None): """ :arg expansion: a subclass of :class:`sympy.expansion.ExpansionBase` :arg strength_usage: A list of integers indicating which expression @@ -54,20 +57,18 @@ class P2EBase(KernelCacheWrapper): number of strength arrays that need to be passed. Default: all kernels use the same strength. """ + if kernels is None: + kernels = [expansion.kernel] - if device is None: - device = ctx.devices[0] + KernelComputation.__init__(self, ctx=ctx, kernels=kernels, + strength_usage=strength_usage, value_dtypes=None, + name=name, device=device) from sumpy.kernel import TargetDerivativeRemover expansion = expansion.with_kernel( TargetDerivativeRemover()(expansion.kernel)) - self.ctx = ctx self.expansion = expansion - self.options = options - self.name = name or self.default_name - self.device = device - self.dim = expansion.dim def get_loopy_instructions(self): @@ -80,24 +81,38 @@ class P2EBase(KernelCacheWrapper): from sumpy.assignment_collection import SymbolicAssignmentCollection sac = SymbolicAssignmentCollection() - coeff_names = [ - sac.assign_unique("coeff%d" % i, coeff_i) + coeff_names = [] + for knl_idx, kernel in enumerate(self.kernels): for i, coeff_i in enumerate( - self.expansion.coefficients_from_source(avec, None, rscale, sac))] + self.expansion.coefficients_from_source(kernel, avec, None, rscale, + sac) + ): + sac.add_assignment(f"coeff{i}_{knl_idx}", coeff_i) + coeff_names.append(f"coeff{i}_{knl_idx}") sac.run_global_cse() + code_transformers = [self.expansion.get_code_transformer()] \ + + [kernel.get_code_transformer() for kernel in self.kernels] + from sumpy.codegen import to_loopy_insns return to_loopy_insns( sac.assignments.items(), vector_names={"a"}, - pymbolic_expr_maps=[self.expansion.get_code_transformer()], + pymbolic_expr_maps=code_transformers, retain_names=coeff_names, complex_dtype=np.complex128 # FIXME ) def get_cache_key(self): - return (type(self).__name__, self.name, self.expansion) + return (type(self).__name__, self.name, self.expansion, tuple(self.kernels), + tuple(self.strength_usage)) + + def get_result_expr(self, icoeff): + isrc = pymbolic.var("isrc") + return sum(pymbolic.var(f"coeff{icoeff}_{i}") + * pymbolic.var("strengths")[self.strength_usage[i], isrc] + for i in range(len(self.kernels))) # }}} @@ -126,21 +141,20 @@ class P2EFromSingleBox(P2EBase): for isrc <> a[idim] = center[idim] - sources[idim, isrc] {dup=idim} - - <> strength = strengths[isrc] """] + self.get_loopy_instructions() + [""" end - """] + [""" + """] + [f""" tgt_expansions[src_ibox-tgt_base_ibox, {coeffidx}] = \ - simul_reduce(sum, isrc, strength*coeff{coeffidx}) \ + simul_reduce(sum, isrc, {self.get_result_expr(coeffidx)}) \ {{id_prefix=write_expn}} - """.format(coeffidx=i) for i in range(ncoeffs)] + [""" + """ for coeffidx in range(ncoeffs)] + [""" end """], [ lp.GlobalArg("sources", None, shape=(self.dim, "nsources"), dim_tags="sep,c"), - lp.GlobalArg("strengths", None, shape="nsources"), + lp.GlobalArg("strengths", None, shape="strength_count, nsources", + dim_tags="sep,C"), lp.GlobalArg("box_source_starts,box_source_counts_nonchild", None, shape=None), lp.GlobalArg("centers", None, shape="dim, aligned_nboxes"), @@ -150,12 +164,13 @@ class P2EFromSingleBox(P2EBase): lp.ValueArg("nboxes,aligned_nboxes,tgt_base_ibox", np.int32), lp.ValueArg("nsources", np.int32), "..." - ] + gather_loopy_source_arguments([self.expansion]), + ] + gather_loopy_source_arguments(self.kernels + (self.expansion,)), name=self.name, assumptions="nsrc_boxes>=1", silenced_warnings="write_race(write_expn*)", default_offset=lp.auto, - fixed_parameters=dict(dim=self.dim), + fixed_parameters=dict(dim=self.dim, + strength_count=self.strength_count), lang_version=MOST_RECENT_LANGUAGE_VERSION) loopy_knl = self.expansion.prepare_loopy_kernel(loopy_knl) @@ -206,7 +221,8 @@ class P2EFromCSR(P2EBase): [ lp.GlobalArg("sources", None, shape=(self.dim, "nsources"), dim_tags="sep,c"), - lp.GlobalArg("strengths", None, shape="nsources"), + lp.GlobalArg("strengths", None, shape="strength_count, nsources", + dim_tags="sep,C"), lp.GlobalArg("source_box_starts,source_box_lists", None, shape=None, offset=lp.auto), lp.GlobalArg("box_source_starts,box_source_counts_nonchild", @@ -218,7 +234,7 @@ class P2EFromCSR(P2EBase): np.int32), lp.ValueArg("nsources", np.int32), "..." - ] + gather_loopy_source_arguments([self.expansion])) + ] + gather_loopy_source_arguments(self.kernels + (self.expansion,))) loopy_knl = lp.make_kernel( [ @@ -244,17 +260,15 @@ class P2EFromCSR(P2EBase): for isrc <> a[idim] = center[idim] - sources[idim, isrc] \ {dup=idim} - - <> strength = strengths[isrc] """] + self.get_loopy_instructions() + [""" end end - """] + [""" + """] + [f""" tgt_expansions[tgt_ibox - tgt_base_ibox, {coeffidx}] = \ simul_reduce(sum, (isrc_box, isrc), - strength*coeff{coeffidx}) \ + {self.get_result_expr(coeffidx)}) \ {{id_prefix=write_expn}} - """.format(coeffidx=i) for i in range(ncoeffs)] + [""" + """ for coeffidx in range(ncoeffs)] + [""" end """], arguments, @@ -262,7 +276,8 @@ class P2EFromCSR(P2EBase): assumptions="ntgt_boxes>=1", silenced_warnings="write_race(write_expn*)", default_offset=lp.auto, - fixed_parameters=dict(dim=self.dim), + fixed_parameters=dict(dim=self.dim, + strength_count=self.strength_count), lang_version=MOST_RECENT_LANGUAGE_VERSION) loopy_knl = self.expansion.prepare_loopy_kernel(loopy_knl) @@ -274,6 +289,7 @@ class P2EFromCSR(P2EBase): # FIXME knl = self.get_kernel() knl = lp.split_iname(knl, "itgt_box", 16, outer_tag="g.0") + return knl def __call__(self, queue, **kwargs): diff --git a/sumpy/p2p.py b/sumpy/p2p.py index e1a8da6d43886c9585ff6a9d5b4f66a27df4e0b2..746a9ef7f9b6bf82881552e6380238b7b1c85008 100644 --- a/sumpy/p2p.py +++ b/sumpy/p2p.py @@ -53,8 +53,7 @@ Particle-to-particle class P2PBase(KernelComputation, KernelCacheWrapper): def __init__(self, ctx, kernels, exclude_self, strength_usage=None, - value_dtypes=None, - options=[], name=None, device=None): + value_dtypes=None, name=None, device=None): """ :arg kernels: list of :class:`sumpy.kernel.Kernel` instances :arg strength_usage: A list of integers indicating which expression @@ -64,7 +63,7 @@ class P2PBase(KernelComputation, KernelCacheWrapper): """ KernelComputation.__init__(self, ctx, kernels, strength_usage, value_dtypes, - name, options, device) + name, device) self.exclude_self = exclude_self diff --git a/sumpy/qbx.py b/sumpy/qbx.py index 7d539587d3d1d2a6725a71b334a5a8bd122b9d4a..9e556ccc7217669c1cb375a452985d846cc78a0f 100644 --- a/sumpy/qbx.py +++ b/sumpy/qbx.py @@ -65,7 +65,8 @@ def stringify_expn_index(i): def expand(expansion_nr, sac, expansion, avec, bvec): rscale = sym.Symbol("rscale") - coefficients = expansion.coefficients_from_source(avec, bvec, rscale) + coefficients = expansion.coefficients_from_source(expansion.kernel, + avec, bvec, rscale) assigned_coeffs = [ sym.Symbol( @@ -85,10 +86,10 @@ def expand(expansion_nr, sac, expansion, avec, bvec): class LayerPotentialBase(KernelComputation, KernelCacheWrapper): def __init__(self, ctx, expansions, strength_usage=None, value_dtypes=None, - options=[], name=None, device=None): + name=None, device=None): KernelComputation.__init__(self, ctx, expansions, strength_usage, value_dtypes, - name, options, device) + name, device) from pytools import single_valued self.dim = single_valued(knl.dim for knl in self.expansions) diff --git a/sumpy/tools.py b/sumpy/tools.py index 98f93cae103643666492e84f70324160e18046d8..1df26c9949e149db2bf14de5045533d5a61f8c4b 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -202,7 +202,7 @@ class KernelComputation: """Common input processing for kernel computations.""" def __init__(self, ctx, kernels, strength_usage, - value_dtypes, name, options=[], device=None): + value_dtypes, name, device=None): """ :arg kernels: list of :class:`sumpy.kernel.Kernel` instances :class:`sumpy.kernel.TargetDerivative` wrappers should be @@ -246,7 +246,7 @@ class KernelComputation: self.context = ctx self.device = device - self.kernels = kernels + self.kernels = tuple(kernels) self.value_dtypes = value_dtypes self.strength_usage = strength_usage self.strength_count = strength_count diff --git a/sumpy/toys.py b/sumpy/toys.py index 6fceebd562c9a3522cc88d7768590bd592ec7f3e..759a765a38e15ee4ff15ce93ae91d7419fcb2afb 100644 --- a/sumpy/toys.py +++ b/sumpy/toys.py @@ -124,35 +124,35 @@ class ToyContext: @memoize_method def get_p2p(self): from sumpy.p2p import P2P - return P2P(self.cl_context, [self.kernel], exclude_self=False) + return P2P(self.cl_context, (self.kernel,), exclude_self=False) @memoize_method def get_p2m(self, order): from sumpy import P2EFromSingleBox return P2EFromSingleBox(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, order), - [self.kernel]) + kernels=(self.kernel,)) @memoize_method def get_p2l(self, order): from sumpy import P2EFromSingleBox return P2EFromSingleBox(self.cl_context, self.local_expn_class(self.no_target_deriv_kernel, order), - [self.kernel]) + kernels=(self.kernel,)) @memoize_method def get_m2p(self, order): from sumpy import E2PFromSingleBox return E2PFromSingleBox(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, order), - [self.kernel]) + (self.kernel,)) @memoize_method def get_l2p(self, order): from sumpy import E2PFromSingleBox return E2PFromSingleBox(self.cl_context, self.local_expn_class(self.no_target_deriv_kernel, order), - [self.kernel]) + (self.kernel,)) @memoize_method def get_m2m(self, from_order, to_order): @@ -197,7 +197,7 @@ def _p2e(psource, center, rscale, order, p2e, expn_class, expn_kwargs): box_source_counts_nonchild=box_source_counts_nonchild, centers=centers, sources=psource.points, - strengths=psource.weights, + strengths=(psource.weights,), rscale=rscale, nboxes=1, tgt_base_ibox=0, diff --git a/test/test_codegen.py b/test/test_codegen.py index 369200001c80cd0e52a3061e5cd3a008980a2643..f580900c5cc4c4e80d83ed1ba17304cb4c9eb717 100644 --- a/test/test_codegen.py +++ b/test/test_codegen.py @@ -89,7 +89,7 @@ def test_line_taylor_coeff_growth(): expn = LineTaylorLocalExpansion(LaplaceKernel(2), order) avec = make_sym_vector("a", 2) bvec = make_sym_vector("b", 2) - coeffs = expn.coefficients_from_source(avec, bvec, rscale=1) + coeffs = expn.coefficients_from_source(expn.kernel, avec, bvec, rscale=1) sym2pymbolic = SympyToPymbolicMapper() coeffs_pymbolic = [sym2pymbolic(c) for c in coeffs] diff --git a/test/test_fmm.py b/test/test_fmm.py index 45c17c38fb42fa6aad21001fed896bb530491ff4..73fd236ce29b7ff8445af9440efa17f4e7ab89ec 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -179,7 +179,7 @@ def test_sumpy_fmm(ctx_factory, knl, local_expn_class, mpole_expn_class): from boxtree.fmm import drive_fmm - pot, = drive_fmm(trav, wrangler, weights) + pot, = drive_fmm(trav, wrangler, (weights,)) from sumpy import P2P p2p = P2P(ctx, out_kernels, exclude_self=False) @@ -249,7 +249,7 @@ def test_sumpy_fmm_timing_data_collection(ctx_factory): from boxtree.fmm import drive_fmm timing_data = {} - pot, = drive_fmm(trav, wrangler, weights, timing_data=timing_data) + pot, = drive_fmm(trav, wrangler, (weights,), timing_data=timing_data) print(timing_data) assert timing_data @@ -308,7 +308,7 @@ def test_sumpy_fmm_exclude_self(ctx_factory): from boxtree.fmm import drive_fmm - pot, = drive_fmm(trav, wrangler, weights) + pot, = drive_fmm(trav, wrangler, (weights,)) from sumpy import P2P p2p = P2P(ctx, out_kernels, exclude_self=True) diff --git a/test/test_kernels.py b/test/test_kernels.py index f8d5122d2ce5ecfce05b0fa85453328127781c6e..3d0dea7a82091fbdf57cccbf88e3b10cd77d3390 100644 --- a/test/test_kernels.py +++ b/test/test_kernels.py @@ -105,6 +105,115 @@ def test_p2p(ctx_factory, exclude_self): assert rel_err < 1e-3 +@pytest.mark.parametrize(("base_knl", "expn_class"), [ + (LaplaceKernel(2), LaplaceConformingVolumeTaylorLocalExpansion), + (LaplaceKernel(2), LaplaceConformingVolumeTaylorMultipoleExpansion), +]) +def test_p2e_multiple(ctx_factory, base_knl, expn_class): + + from sympy.core.cache import clear_cache + clear_cache() + + order = 4 + ctx = ctx_factory() + queue = cl.CommandQueue(ctx) + + np.random.seed(17) + + nsources = 100 + + extra_kwargs = {} + if isinstance(base_knl, HelmholtzKernel): + if base_knl.allow_evanescent: + extra_kwargs["k"] = 0.2 * (0.707 + 0.707j) + else: + extra_kwargs["k"] = 0.2 + if isinstance(base_knl, StokesletKernel): + extra_kwargs["mu"] = 0.2 + + in_kernels = [ + DirectionalSourceDerivative(base_knl, "dir_vec"), + base_knl, + ] + + knl = base_knl + expn = expn_class(knl, order=order) + + from sumpy import P2EFromSingleBox + + center = np.array([2, 1, 0][:knl.dim], np.float64) + sources = (0.7*(-0.5+np.random.rand(knl.dim, nsources).astype(np.float64)) + + center[:, np.newaxis]) + + strengths = [ + np.ones(nsources, dtype=np.float64) * (1/nsources), + np.ones(nsources, dtype=np.float64) * (2/nsources) + ] + + source_boxes = np.array([0], dtype=np.int32) + box_source_starts = np.array([0], dtype=np.int32) + box_source_counts_nonchild = np.array([nsources], dtype=np.int32) + + alpha = np.linspace(0, 2*np.pi, nsources, np.float64) + dir_vec = np.vstack([np.cos(alpha), np.sin(alpha)]) + + from sumpy.expansion.local import LocalExpansionBase + 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) + else: + centers = (np.array([0.0, 0.0, 0.0][:knl.dim], + dtype=np.float64).reshape(knl.dim, 1) + + center[:, np.newaxis]) + + rscale = 0.5 # pick something non-1 + + # apply p2e at the same time + p2e = P2EFromSingleBox(ctx, expn, kernels=in_kernels, strength_usage=[0, 1]) + evt, (mpoles,) = p2e(queue, + source_boxes=source_boxes, + box_source_starts=box_source_starts, + box_source_counts_nonchild=box_source_counts_nonchild, + centers=centers, + sources=sources, + strengths=strengths, + nboxes=1, + tgt_base_ibox=0, + rscale=rscale, + + #flags="print_hl_cl", + out_host=True, + dir_vec=dir_vec, + **extra_kwargs) + + actual_result = mpoles + + # apply p2e separately + expected_result = np.zeros_like(actual_result) + for i, in_kernel in enumerate(in_kernels): + extra_source_kwargs = extra_kwargs.copy() + if isinstance(in_kernel, DirectionalSourceDerivative): + extra_source_kwargs["dir_vec"] = dir_vec + p2e = P2EFromSingleBox(ctx, expn, kernels=[in_kernel], strength_usage=[i]) + evt, (mpoles,) = p2e(queue, + source_boxes=source_boxes, + box_source_starts=box_source_starts, + box_source_counts_nonchild=box_source_counts_nonchild, + centers=centers, + sources=sources, + strengths=strengths, + nboxes=1, + tgt_base_ibox=0, + rscale=rscale, + + #flags="print_hl_cl", + out_host=True, **extra_source_kwargs) + expected_result += mpoles + + norm = la.norm(actual_result - expected_result)/la.norm(expected_result) + assert norm < 1e-12 + + @pytest.mark.parametrize("order", [4]) @pytest.mark.parametrize(("base_knl", "expn_class"), [ (LaplaceKernel(2), VolumeTaylorLocalExpansion), @@ -173,8 +282,8 @@ def test_p2e2p(ctx_factory, base_knl, expn_class, order, with_source_derivative) expn = expn_class(knl, order=order) from sumpy import P2EFromSingleBox, E2PFromSingleBox, P2P - p2e = P2EFromSingleBox(ctx, expn, out_kernels) - e2p = E2PFromSingleBox(ctx, expn, out_kernels) + p2e = P2EFromSingleBox(ctx, expn, kernels=[knl]) + e2p = E2PFromSingleBox(ctx, expn, kernels=out_kernels) p2p = P2P(ctx, out_kernels, exclude_self=False) from pytools.convergence import EOCRecorder @@ -229,7 +338,7 @@ def test_p2e2p(ctx_factory, base_knl, expn_class, order, with_source_derivative) box_source_counts_nonchild=box_source_counts_nonchild, centers=centers, sources=sources, - strengths=strengths, + strengths=(strengths,), nboxes=1, tgt_base_ibox=0, rscale=rscale, @@ -495,7 +604,7 @@ def test_translations(ctx_factory, knl, local_expn_class, mpole_expn_class): box_source_counts_nonchild=p2m_box_source_counts_nonchild, centers=centers, sources=sources, - strengths=strengths, + strengths=(strengths,), nboxes=nboxes, rscale=m1_rscale,