diff --git a/sumpy/kernel.py b/sumpy/kernel.py index a68659268d2432b79b51c9b7c6c0a4e1f430a6e8..75500808ea4d21af6af3d60132623d145bf986c9 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -492,11 +492,10 @@ class StokesletKernel(ExpressionKernel): class StressletKernel(ExpressionKernel): - init_arg_names = ("dim", "icomp", "jcomp", "viscosity_mu_name", - "stresslet_vector_name") + init_arg_names = ("dim", "icomp", "jcomp", "kcomp", "viscosity_mu_name") - def __init__(self, dim=None, icomp=None, jcomp=None, viscosity_mu_name="mu", - stresslet_vector_name="stresslet_vec"): + def __init__(self, dim=None, icomp=None, jcomp=None, kcomp=None, + viscosity_mu_name="mu"): """ :arg viscosity_mu_name: The argument name to use for dynamic viscosity :math:`\mu` the then generating functions to @@ -506,23 +505,17 @@ 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 + d[icomp]*d[jcomp]*d[kcomp]/r**4 ) scaling = 1/(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 + d[icomp]*d[jcomp]*d[kcomp]/r**5 ) scaling = -3/(4*var("pi")) @@ -536,6 +529,7 @@ class StressletKernel(ExpressionKernel): self.stresslet_vector_name = stresslet_vector_name self.icomp = icomp self.jcomp = jcomp + self.kcomp = kcomp ExpressionKernel.__init__( self, @@ -545,42 +539,26 @@ class StressletKernel(ExpressionKernel): is_complex_valued=False) def __getinitargs__(self): - return (self._dim, self.icomp, self.jcomp, self.viscosity_mu_name, - self.stresslet_vector_name) + return (self._dim, self.icomp, self.jcomp, self.kcomp, + self.viscosity_mu_name) def update_persistent_hash(self, key_hash, key_builder): key_hash.update(type(self).__name__.encode()) key_builder.rec(key_hash, ( - self.dim, self.icomp, self.jcomp, self.viscosity_mu_name, + self.dim, self.icomp, self.jcomp, self.kcomp, self.viscosity_mu_name, self.stresslet_vector_name)) def __repr__(self): - return "StressletKnl%dD_%d%d[%s]" % (self.dim, self.icomp, self.jcomp, - self.stresslet_vector_name) + return "StressletKnl%dD_%d%d%d" % (self.dim, self.icomp, self.jcomp, + self.kcomp) def get_args(self): return [ KernelArgument( loopy_arg=lp.ValueArg(self.viscosity_mu_name, np.float64), - ), - KernelArgument( - loopy_arg=lp.GlobalArg(self.stresslet_vector_name, - None, - shape=(self.dim, "nsources"), - dim_tags="sep,C")) + ) ] - def get_code_transformer(self): - from sumpy.codegen import VectorComponentRewriter - vcr = VectorComponentRewriter([self.stresslet_vector_name]) - from pymbolic.primitives import Variable - via = _VectorIndexAdder(self.stresslet_vector_name, (Variable("isrc"),)) - - def transform(expr): - return via(vcr(expr)) - - return transform - mapper_method = "map_stresslet_kernel" # }}}