diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 584b2240543ff99a1c72377f1109e8c2256790f3..2f5c3356fa094b25999dd59cafb2798793f77ef9 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -788,7 +788,12 @@ class StressletKernel(ExpressionKernel): dynamic viscosity :math:`\mu` the then generating functions to evaluate this kernel. """ - # Mu is unused but kept for consistency with the stokeslet. + # mu is unused but kept for consistency with the Stokeslet. + if isinstance(viscosity_mu, str): + mu = parse(viscosity_mu) + else: + mu = viscosity_mu + if dim == 2: d = make_sym_vector("d", dim) r = pymbolic_real_norm_2(d) @@ -811,7 +816,7 @@ class StressletKernel(ExpressionKernel): self.icomp = icomp self.jcomp = jcomp self.kcomp = kcomp - self.viscosity_mu = viscosity_mu + self.viscosity_mu = mu super().__init__( dim, @@ -824,13 +829,24 @@ class StressletKernel(ExpressionKernel): 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.kcomp, self.viscosity_mu)) + key_builder.rec(key_hash, (self.dim, self.icomp, self.jcomp, self.kcomp)) + + from pymbolic.mapper.persistent_hash import PersistentHashWalkMapper + mapper = PersistentHashWalkMapper(key_hash) + mapper(self.viscosity_mu) def __repr__(self): return "StressletKnl%dD_%d%d%d" % (self.dim, self.icomp, self.jcomp, self.kcomp) + @memoize_method + def get_args(self): + from sumpy.tools import get_all_variables + variables = get_all_variables(self.viscosity_mu) + return [ + KernelArgument(loopy_arg=lp.ValueArg(v.name, np.float64)) + for v in variables] + mapper_method = "map_stresslet_kernel" def get_pde_as_diff_op(self): @@ -838,13 +854,6 @@ class StressletKernel(ExpressionKernel): w = make_identity_diff_op(self.dim) return laplacian(laplacian(w)) - def get_args(self): - return [ - KernelArgument( - loopy_arg=lp.ValueArg(self.viscosity_mu, np.float64), - ) - ] - class LineOfCompressionKernel(ExpressionKernel): """A kernel for the line of compression or dilatation of constant strength