diff --git a/sumpy/codegen.py b/sumpy/codegen.py index bef7fc6bf93a171a96626f9db88c312c23167bf2..02da88ff52d2d227c4203ef273ec2ab4a7b7d78a 100644 --- a/sumpy/codegen.py +++ b/sumpy/codegen.py @@ -699,6 +699,25 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[], # cse_walk(expr) #cse_tag = CSETagMapper(cse_walk) + # {{{ complex constant mapper + + class ComplexConstantMapper(IdentityMapper): + """Map complex constants to, for example, add additional size info. + """ + def __init__(self, complex_dtype=None): + if complex_dtype is None: + self.complex_dtype = complex + else: + self.complex_dtype = complex_dtype + + def map_constant(self, expr, *args, **kwargs): + if np.asarray(expr).dtype.kind == 'c': + return self.complex_dtype(expr) + else: + return expr + + # }}} End complex constant mapper + # do the rest of the conversion bessel_sub = BesselSubstitutor(BesselGetter(btog.bessel_j_arg_to_top_order)) vcr = VectorComponentRewriter(vector_names) @@ -706,6 +725,7 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[], ssg = SumSignGrouper() fck = FractionKiller() bik = BigIntegerKiller() + ccm = ComplexConstantMapper(complex_dtype) cmr = ComplexRewriter() def convert_expr(name, expr): @@ -717,6 +737,7 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[], expr = fck(expr) expr = ssg(expr) expr = bik(expr) + expr = ccm(expr) expr = cmr(expr) #expr = cse_tag(expr) for m in pymbolic_expr_maps: diff --git a/sumpy/e2p.py b/sumpy/e2p.py index 7b0072ad55708ba205ceed3448914ad0d8eda281..214a1e74e95a5643afb348eb5df366d5a79ac127 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -32,6 +32,8 @@ import sumpy.symbolic as sym from sumpy.tools import KernelCacheWrapper from loopy.version import MOST_RECENT_LANGUAGE_VERSION +import logging +logger = logging.getLogger(__name__) __doc__ = """ @@ -320,7 +322,13 @@ class E2PFromCSR(E2PBase): # meaningfully inferred. Make the type of rscale explicit. rscale = centers.dtype.type(kwargs.pop("rscale")) - return knl(queue, centers=centers, rscale=rscale, **kwargs) + try: + return knl(queue, centers=centers, rscale=rscale, **kwargs) + except TypeError: + logger.debug("Initial trial of p2e failed, " + "probably the loopy kernel does not accept rscale") + logger.debug("Trying to make self-adjustments") + return knl(queue, centers=centers, **kwargs) # }}} diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 3936c518246aa8f14e54dba2f0b33d5361da04c6..8e25a7fb67bcdb533245ce347c176c3831cb384b 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -64,6 +64,9 @@ of them in the process. .. autoclass:: AxisTargetDerivative .. autoclass:: DirectionalTargetDerivative .. autoclass:: DirectionalSourceDerivative +.. autoclass:: LaplacianDerivativeBase +.. autoclass:: LaplacianTargetDerivative +.. autoclass:: LaplacianSourceDerivative Transforming kernels -------------------- @@ -736,6 +739,70 @@ class StressletKernel(ExpressionKernel): mapper_method = "map_stresslet_kernel" + +class FactorizedBiharmonicKernel(ExpressionKernel): + init_arg_names = ("dim", "lambda1_name", "lambda2_name") + + def __init__(self, dim=None, lambda1_name="lam1", lambda2_name="lam2"): + + lam = [var(lambda1_name), var(lambda2_name)] + + if dim == 2: + r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + + # http://dlmf.nist.gov/10.27#E8 + expr_for_K0 = [var("hankel_1")(0, var("I")*lam[i]*r) # noqa: N806 + for i in range(2)] + scaling_for_K0 = 1/2*var("pi")*var("I") # noqa: N806 + + expr = expr_for_K0[0] - expr_for_K0[1] + scaling = -1/(2 * var("pi") * ( + lam[0]**2 - lam[1]**2)) * scaling_for_K0 + else: + raise NotImplementedError("unsupported dimensionality") + + super(FactorizedBiharmonicKernel, self).__init__( + dim, + expression=expr, + global_scaling_const=scaling, + is_complex_valued=True) + + self.lambda1_name = lambda1_name + self.lambda2_name = lambda2_name + + def __getinitargs__(self): + return(self.dim, self.lambda1_name, self.lambda2_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.lambda1_name, self.lambda2_name)) + + def __repr__(self): + return "FctBiharmKnl%dD(%s,%s)" % ( + self.dim, self.lambda1_name, self.lambda2_name) + + def prepare_loopy_kernel(self, loopy_knl): + from sumpy.codegen import (bessel_preamble_generator, bessel_mangler) + loopy_knl = lp.register_function_manglers(loopy_knl, + [bessel_mangler]) + loopy_knl = lp.register_preamble_generators(loopy_knl, + [bessel_preamble_generator]) + + return loopy_knl + + def get_args(self): + return [ + KernelArgument( + loopy_arg=lp.ValueArg(self.lambda1_name, np.float64), + ), + KernelArgument( + loopy_arg=lp.ValueArg(self.lambda2_name, np.float64) + ) + ] + + mapper_method = "map_factorized_biharmonic_kernel" + # }}} @@ -952,6 +1019,81 @@ class DirectionalSourceDerivative(DirectionalDerivative): mapper_method = "map_directional_source_derivative" + +class LaplacianDerivativeBase(DerivativeBase): + init_arg_names = ("inner_kernel", ) + + def __init__(self, inner_kernel): + KernelWrapper.__init__(self, inner_kernel) + + def __getinitargs__(self): + return (self.inner_kernel, ) + + def update_persistent_hash(self, key_hash, key_builder): + key_hash.update(type(self).__name__.encode("utf8")) + key_builder.rec(key_hash, self.inner_kernel) + + def __str__(self): + return r"Lap_%s %s" % ( + self.laplacian_kind[0], self.inner_kernel) + + def __repr__(self): + return "%s(%s)" % ( + type(self).__name__, + self.inner_kernel) + + def get_source_args(self): + return self.inner_kernel.get_source_args() + + +class LaplacianTargetDerivative(LaplacianDerivativeBase): + laplacian_kind = "tgt" + + def get_code_transformer(self): + inner = self.inner_kernel.get_code_transformer() + + def transform(expr): + return inner(expr) + + return transform + + def postprocess_at_target(self, expr, bvec): + expr = self.inner_kernel.postprocess_at_target(expr, bvec) + + dimensions = len(bvec) + assert dimensions == self.dim + + # vec_r = target - bvec + return sum(expr.diff(bvec[axis]).diff(bvec[axis]) + for axis in range(dimensions)) + + mapper_method = "map_laplacian_target_derivative" + + +class LaplacianSourceDerivative(LaplacianDerivativeBase): + laplacian_kind = "src" + + def get_code_transformer(self): + inner = self.inner_kernel.get_code_transformer() + + def transform(expr): + return inner(expr) + + return transform + + def postprocess_at_source(self, expr, avec): + expr = self.inner_kernel.postprocess_at_source(expr, avec) + + dimensions = len(avec) + assert dimensions == self.dim + + # vec_r = avec - source, (-1)^2 cancels out + return sum(expr.diff(avec[axis]).diff(avec[axis]) + for axis in range(dimensions)) + + mapper_method = "map_laplacian_source_derivative" + + # }}} @@ -993,6 +1135,7 @@ class KernelIdentityMapper(KernelMapper): map_yukawa_kernel = map_expression_kernel map_stokeslet_kernel = map_expression_kernel map_stresslet_kernel = map_expression_kernel + map_factorized_biharmonic_kernel = map_expression_kernel def map_axis_target_derivative(self, kernel): return AxisTargetDerivative(kernel.axis, self.rec(kernel.inner_kernel)) @@ -1004,6 +1147,11 @@ class KernelIdentityMapper(KernelMapper): map_directional_source_derivative = map_directional_target_derivative + def map_laplacian_target_derivative(self, kernel): + return type(kernel)(self.rec(kernel.inner_kernel)) + + map_laplacian_source_derivative = map_laplacian_target_derivative + class AxisTargetDerivativeRemover(KernelIdentityMapper): def map_axis_target_derivative(self, kernel): @@ -1014,11 +1162,17 @@ class TargetDerivativeRemover(AxisTargetDerivativeRemover): def map_directional_target_derivative(self, kernel): return self.rec(kernel.inner_kernel) + def map_laplacian_target_derivative(self, kernel): + return self.rec(kernel.inner_kernel) + class SourceDerivativeRemover(KernelIdentityMapper): def map_directional_source_derivative(self, kernel): return self.rec(kernel.inner_kernel) + def map_laplacian_source_derivative(self, kernel): + return self.rec(kernel.inner_kernel) + class DerivativeCounter(KernelCombineMapper): def combine(self, values): @@ -1040,6 +1194,11 @@ class DerivativeCounter(KernelCombineMapper): map_directional_target_derivative = map_axis_target_derivative map_directional_source_derivative = map_axis_target_derivative + def map_laplacian_target_derivative(self, kernel): + return 2 + self.rec(kernel.inner_kernel) + + map_laplacian_source_derivative = map_laplacian_target_derivative + # }}} diff --git a/sumpy/p2e.py b/sumpy/p2e.py index daa7d93dc5b393d358424031500a4915e2c2dd9f..db8f491ecd473b9da9a256db7e95c8311fb1464e 100644 --- a/sumpy/p2e.py +++ b/sumpy/p2e.py @@ -299,7 +299,13 @@ class P2EFromCSR(P2EBase): # meaningfully inferred. Make the type of rscale explicit. rscale = centers.dtype.type(kwargs.pop("rscale")) - return knl(queue, centers=centers, rscale=rscale, **kwargs) + try: + return knl(queue, centers=centers, rscale=rscale, **kwargs) + except TypeError: + logger.debug("Initial trial of p2e failed, " + "probably the loopy kernel does not accept rscale") + logger.debug("Trying to make self-adjustments") + return knl(queue, centers=centers, **kwargs) # }}}