diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 9c8a4958713800d37c18babd57f2426340a36222..0cb14028261ca7d8a977c26f340c5606f6fd26e3 100644 --- a/sumpy/codegen.py +++ b/sumpy/codegen.py @@ -430,6 +430,41 @@ class FractionKiller(CSECachingMapperMixin, IdentityMapper): # }}} +# {{{ convert big integers into floats + +from loopy.tools import is_integer + + +class BigIntegerKiller(CSECachingMapperMixin, IdentityMapper): + + def __init__(self, warn_on_first=True, int_type=np.int64, float_type=np.float64): + IdentityMapper.__init__(self) + self.bits = 64 + self.warn = warn_on_first + self.float_type = float_type + self.iinfo = np.iinfo(int_type) + + def map_constant(self, expr): + if not is_integer(expr): + return IdentityMapper.map_constant(self, expr) + + if self.iinfo.min <= expr <= self.iinfo.max: + return expr + + if self.warn: + from warnings import warn + warn("Converting '%d' to float: this may result in " + "loss of digits." % expr) + # Suppress further warnings. + self.warn = False + + return self.float_type(expr) + + map_common_subexpression_uncached = IdentityMapper.map_common_subexpression + +# }}} + + # {{{ vector component rewriter INDEXED_VAR_RE = re.compile("^([a-zA-Z_]+)([0-9]+)$") @@ -534,6 +569,7 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[], pwr = PowerRewriter() ssg = SumSignGrouper() fck = FractionKiller() + bik = BigIntegerKiller() def convert_expr(name, expr): logger.debug("generate expression for: %s" % name) @@ -543,6 +579,7 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[], expr = pwr(expr) expr = fck(expr) expr = ssg(expr) + expr = bik(expr) #expr = cse_tag(expr) for m in pymbolic_expr_maps: expr = m(expr) diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 39f490a3f765dae50c9cce2e0be47f564db31796..4a3987b5e03893f32510ae2404a7cee71b5e7639 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -506,26 +506,22 @@ class StressletKernel(ExpressionKernel): if dim == 2: d = make_sym_vector("d", dim) - n = make_sym_vector(stresslet_vector_name, dim) r = pymbolic_real_norm_2(d) expr = ( - sum(n[axis]*d[axis] for axis in range(dim)) - * - d[icomp]*d[jcomp]/r**4 + -var("log")(r)*(1 if icomp == jcomp else 0) + + + d[icomp]*d[jcomp]/r**2 ) - scaling = 1/(var("pi")) - + scaling = 1/(4*var("pi")) elif dim == 3: d = make_sym_vector("d", dim) - n = make_sym_vector(stresslet_vector_name, dim) r = pymbolic_real_norm_2(d) expr = ( - sum(n[axis]*d[axis] for axis in range(dim)) - * - d[icomp]*d[jcomp]/r**5 + (1/r)*(1 if icomp == jcomp else 0) + + + d[icomp]*d[jcomp]/r**3 ) - scaling = -3/(4*var("pi")) - + scaling = -1/(8*var("pi")) elif dim is None: expr = None scaling = None @@ -562,15 +558,28 @@ class StressletKernel(ExpressionKernel): return [ KernelArgument( loopy_arg=lp.ValueArg(self.viscosity_mu_name, np.float64), - ), + )] + + def get_source_args(self): + return [ KernelArgument( loopy_arg=lp.GlobalArg(self.stresslet_vector_name, None, shape=(self.dim, "nsources"), - dim_tags="sep,C")) - ] + dim_tags="sep,C"))] + + def postprocess_at_source(self, expr, avec): + dimensions = len(avec) + assert dimensions == self.dim - def get_code_tranformer(self): + from sumpy.symbolic import make_sympy_vector + n = make_sympy_vector(self.stresslet_vector_name, dimensions) + + # avec = center-src -> minus sign from chain rule + return sum(-n[axis]*expr.diff(avec[axis]) + for axis in range(dimensions)) + + def get_code_transformer(self): from sumpy.codegen import VectorComponentRewriter vcr = VectorComponentRewriter([self.stresslet_vector_name]) from pymbolic.primitives import Variable diff --git a/sumpy/p2p.py b/sumpy/p2p.py index 7234180f62ff12d63b2e7c5afdaeaecf12c93f87..0345f9046c1e97a43b59cc84076ae63c886b2bc9 100644 --- a/sumpy/p2p.py +++ b/sumpy/p2p.py @@ -154,8 +154,8 @@ class P2P(P2PBase): shape=(self.dim, "nsources")), lp.GlobalArg("targets", None, shape=(self.dim, "ntargets")), - lp.ValueArg("nsources", None), - lp.ValueArg("ntargets", None), + lp.ValueArg("nsources", np.int32), + lp.ValueArg("ntargets", np.int32), lp.GlobalArg("strength", None, shape="nstrengths,nsources"), lp.GlobalArg("result", None, shape="nresults,ntargets", dim_tags="sep,C") diff --git a/sumpy/qbx.py b/sumpy/qbx.py index 0b3dc03537e788a75412b712e1c3ad18b14c11db..629f9092bb88a564a0c822d1f844a9326ddbab49 100644 --- a/sumpy/qbx.py +++ b/sumpy/qbx.py @@ -111,8 +111,8 @@ class LayerPotentialBase(KernelComputation, KernelCacheWrapper): shape=(self.dim, "ntargets"), order="C"), lp.GlobalArg("center", None, shape=(self.dim, "ntargets"), order="C"), - lp.ValueArg("nsources", None), - lp.ValueArg("ntargets", None), + lp.ValueArg("nsources", np.int32), + lp.ValueArg("ntargets", np.int32), ] @memoize_method