diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 7324a082f166b21bc46aa4b95fa6be9dfbb2168e..cf0835de455ec0aecbc03a809aa8bae5a28c9160 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -26,7 +26,7 @@ from six.moves import range, zip import loopy as lp import numpy as np -from pymbolic.mapper import IdentityMapper +from pymbolic.mapper import IdentityMapper, CSECachingMapperMixin from sumpy.symbolic import pymbolic_real_norm_2 from pymbolic.primitives import make_sym_vector from pymbolic import var @@ -652,7 +652,7 @@ class AxisTargetDerivative(DerivativeBase): mapper_method = "map_axis_target_derivative" -class _VectorIndexAdder(IdentityMapper): +class _VectorIndexAdder(CSECachingMapperMixin, IdentityMapper): def __init__(self, vec_name, additional_indices): self.vec_name = vec_name self.additional_indices = additional_indices @@ -666,6 +666,8 @@ class _VectorIndexAdder(IdentityMapper): else: return IdentityMapper.map_subscript(self, expr) + map_common_subexpression_uncached = IdentityMapper.map_common_subexpression + class DirectionalDerivative(DerivativeBase): init_arg_names = ("inner_kernel", "dir_vec_name")