diff --git a/sumpy/e2p.py b/sumpy/e2p.py index 2a4ff531c0522514bb43cbb8267bc411ea9e05fa..3033f0c43826a2ad52cd9e9437ff359b772fe94e 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -57,15 +57,16 @@ class E2PBase(KernelCacheWrapper): device = ctx.devices[0] from sumpy.kernel import (SourceTransformationRemover, - TargetTransformationRemover) + TargetTransformationRemover, AsymptoticScalingRemover) sxr = SourceTransformationRemover() txr = TargetTransformationRemover() + adr = AsymptoticScalingRemover() expansion = expansion.with_kernel( - sxr(expansion.kernel)) + adr(sxr(expansion.kernel))) kernels = [sxr(knl) for knl in kernels] for knl in kernels: - assert txr(knl) == expansion.kernel + assert adr(txr(knl)) == expansion.kernel self.ctx = ctx self.expansion = expansion @@ -74,8 +75,9 @@ class E2PBase(KernelCacheWrapper): self.device = device self.dim = expansion.dim + self._vector_names = {"b"} - def get_loopy_insns_and_result_names(self): + def get_loopy_insns_and_result_names(self, use_qbmax=False): from sumpy.symbolic import make_sym_vector bvec = make_sym_vector("b", self.dim) @@ -90,7 +92,8 @@ class E2PBase(KernelCacheWrapper): result_names = [ sac.assign_unique("result_%d_p" % i, - self.expansion.evaluate(knl, coeff_exprs, bvec, rscale, sac=sac)) + self.expansion.evaluate(knl, coeff_exprs, bvec, rscale, sac=sac, + use_qbmax=use_qbmax)) for i, knl in enumerate(self.kernels) ] @@ -99,7 +102,7 @@ class E2PBase(KernelCacheWrapper): from sumpy.codegen import to_loopy_insns loopy_insns = to_loopy_insns( sac.assignments.items(), - vector_names={"b"}, + vector_names=self._vector_names, pymbolic_expr_maps=[ knl.get_code_transformer() for knl in self.kernels], retain_names=result_names, diff --git a/sumpy/expansion/__init__.py b/sumpy/expansion/__init__.py index 8bfb85e9670375a352a3ee48255b1ddc5a9e1f1b..41071aaf2326d73b94f41daa98b6733ffa128e10 100644 --- a/sumpy/expansion/__init__.py +++ b/sumpy/expansion/__init__.py @@ -141,7 +141,7 @@ class ExpansionBase: result[i] += weight * coeffs[i] return result - def evaluate(self, kernel, coeffs, bvec, rscale, sac=None): + def evaluate(self, kernel, coeffs, bvec, rscale, sac=None, use_qbmax=False): """ :return: a :mod:`sympy` expression corresponding to the evaluated expansion with the coefficients diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py index 726ab0863c81a34c05333cf341c3b0197ee3518e..9cc9d3c89a0787ee43f62c8eadc83e74a6590550 100644 --- a/sumpy/expansion/local.py +++ b/sumpy/expansion/local.py @@ -223,7 +223,7 @@ class LineTaylorLocalExpansion(LocalExpansionBase): .subs("tau", 0) for i in self.get_coefficient_identifiers()] - def evaluate(self, tgt_kernel, coeffs, bvec, rscale, sac=None): + def evaluate(self, tgt_kernel, coeffs, bvec, rscale, sac=None, use_qbmax=False): # no point in heeding rscale here--just ignore it from pytools import factorial return sym.Add(*( @@ -258,34 +258,67 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): """ if not self.use_rscale: rscale = 1 - base_kernel = single_valued(knl.get_base_kernel() for knl in kernels) - base_taker = base_kernel.get_derivative_taker(avec, rscale, sac) - result = [0]*len(self) + # TODO: use a kernel mapper to find asymptotic info wrapper that is not + # the outmost layer + from sumpy.kernel import AsymptoticScalingGetter + asg = AsymptoticScalingGetter() + + # bvec is None when + # 1) when using FMM without QBMAX + # 2) when using FMM with QBMAX, but during GIGAQBX's far-field P2L + using_qbmax = (asg(base_kernel) is not None) and (bvec is not None) + + if using_qbmax: + # With a directional expansion (like line taylor), we write: + # source to target = ( + # source to center (avec) + center to target (bvec)). + # Kernel evaluated wrt rad = avec + bvec, postprocessed wrt avec, + # then take derivatives along the direction of bvec, and + # finally let bvec=0 + base_taker = base_kernel.get_expression(avec + bvec) + post_diff_subs = {bvc: 0 for bvc in bvec} + else: + base_taker = base_kernel.get_derivative_taker(avec, rscale, sac) + post_diff_subs = dict() + # to signal postprocess_at_source that we are not using qbmax + bvec = None + result = [0]*len(self) for knl, weight in zip(kernels, weights): - taker = knl.postprocess_at_source(base_taker, avec) + # apply source derivatives, if any + if using_qbmax: + # Since DifferentiatedExprDerivativeTaker does not work with + # anisotropic expansions, using vanilla expression here is + # necessary. That is, the base_taker is directly constructed + # from \partial_{n(y)}G(x-y) + from sumpy.tools import ExprDerivativeTaker + ppkernel = knl.postprocess_at_source(base_taker, avec, bvec) + base_taker = ExprDerivativeTaker(ppkernel, bvec, rscale, sac) + + taker = knl.postprocess_at_source(base_taker, avec, bvec) + # Following is a hack to make sure cse works. if 1: def save_temp(x): return add_to_sac(sac, weight * x) for i, mi in enumerate(self.get_coefficient_identifiers()): - result[i] += taker.diff(mi, save_temp) + result[i] += taker.diff(mi, save_temp, subs=post_diff_subs) else: def save_temp(x): return add_to_sac(sac, x) for i, mi in enumerate(self.get_coefficient_identifiers()): - result[i] += weight * taker.diff(mi, save_temp) - + result[i] += weight * taker.diff( + mi, save_temp, subs=post_diff_subs) return result def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None): return self.coefficients_from_source_vec((kernel,), avec, bvec, rscale=rscale, weights=(1,), sac=sac) - def evaluate(self, kernel, coeffs, bvec, rscale, sac=None): + def evaluate(self, kernel, coeffs, bvec, rscale, sac=None, use_qbmax=False): if not self.use_rscale: rscale = 1 @@ -296,14 +329,74 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): bvec_scaled = [b*rscale**-1 for b in bvec] from sumpy.tools import mi_power, mi_factorial - result = sum( - coeff - * mi_power(bvec_scaled, mi, evaluate=False) - / mi_factorial(mi) - for coeff, mi in zip( + from sumpy.kernel import AsymptoticScalingGetter, DerivativeCounter + asg = AsymptoticScalingGetter() + dcr = DerivativeCounter() + scaling_expr = asg(kernel) + nder = dcr(kernel) + # When using FMM, we want to turn off qbmax for the far field, + # thus `use_qbmax` must be explicitly set to `True` for qbmax. + if use_qbmax and (scaling_expr is not None) and (nder > 0): + # Target derivatives for qbmax require postprocessing with + # the Leibniz formula. + # + # NOTE: only first order target derivative is implemented. + # The kernel passed to this method during QBX should only have + # target derivatives. + if nder > 1: + raise NotImplementedError + + # Potential = (int[f(x)*G(x,y)dx] * g(y) * g^{-1}(y))' + # = Expand{int[f(x)*G(x,y)dx] * g(y)}' / g(y) + # - Expand{int[f(x)*G(x,y)dx] * g(y)} * g'(y)/g^2(y) + #:= sum A_n(y) / g(y) - sum B_n(y) * g'(y)/g^2(y) + # = sum[A_n(y)/g(y) - B_n(y)*g'(y)/g^2(y)] + # = sum[A_n(y) - B_n(y)*g'(y)/g(y)] / g(y) + # ^: A_n are coeffs of usual target derivative + # ^: B_n are original coeffs + + # get the map coefficients at the target + expr_dict = {(0,)*self.dim: 1} + expr_dict = kernel.get_derivative_transformation_at_target(expr_dict) + + def orig_term(coeff, mi): # sympy expr for B_mi + # B_mi = (taylor_coeff * rscale^mi) * (((x-center)/rscale)^mi/mi!) + # ^: coeff + return (coeff * mi_power(bvec_scaled, mi, evaluate=False) + / mi_factorial(mi)) + + def tder_term(coeff, mi): # sympy expr for A_mi + expr = orig_term(coeff, mi) + return sum( + dcoeff * kernel._diff(expr, bvec, mi) + for mi, dcoeff in expr_dict.items()) + + # get sympy expr for g'/g + g = 1 / scaling_expr + gpog = sum( + coeff * kernel._diff(g, bvec, mi) / g + for mi, coeff in expr_dict.items()) + + def leibniz_term(coeff, mi): + ami = tder_term(coeff, mi) + bmi = orig_term(coeff, mi) + return ami - bmi * gpog + + result = sum( + leibniz_term(coeff, mi) + for coeff, mi in zip( evaluated_coeffs, self.get_full_coefficient_identifiers())) - return kernel.postprocess_at_target(result, bvec) + else: + result = sum( + coeff + * mi_power(bvec_scaled, mi, evaluate=False) + / mi_factorial(mi) + for coeff, mi in zip( + evaluated_coeffs, self.get_full_coefficient_identifiers())) + result = kernel.postprocess_at_target(result, bvec) + + return result def m2l_translation_classes_dependent_ndata(self, src_expansion): """Returns number of expressions in M2L global precomputation step. @@ -783,7 +876,7 @@ class _FourierBesselLocalExpansion(LocalExpansionBase): * sym.exp(sym.I * c * source_angle_rel_center), avec) for c in self.get_coefficient_identifiers()] - def evaluate(self, kernel, coeffs, bvec, rscale, sac=None): + def evaluate(self, kernel, coeffs, bvec, rscale, sac=None, use_qbmax=False): if not self.use_rscale: rscale = 1 diff --git a/sumpy/expansion/multipole.py b/sumpy/expansion/multipole.py index 408b90a9c720aecd7950570544641462afb398f2..4cdc0e715521ee11b059d0b6cc96aa9418ec812d 100644 --- a/sumpy/expansion/multipole.py +++ b/sumpy/expansion/multipole.py @@ -85,7 +85,7 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): return self.coefficients_from_source_vec((kernel,), avec, bvec, rscale, (1,), sac=sac) - def evaluate(self, kernel, coeffs, bvec, rscale, sac=None): + def evaluate(self, kernel, coeffs, bvec, rscale, sac=None, use_qbmax=False): if not self.use_rscale: rscale = 1 @@ -420,7 +420,7 @@ class _HankelBased2DMultipoleExpansion(MultipoleExpansionBase): avec) for c in self.get_coefficient_identifiers()] - def evaluate(self, kernel, coeffs, bvec, rscale, sac=None): + def evaluate(self, kernel, coeffs, bvec, rscale, sac=None, use_qbmax=False): if not self.use_rscale: rscale = 1 diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 584b2240543ff99a1c72377f1109e8c2256790f3..fb58346b74e3d883225a1c9c63ba387712c45d05 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -23,9 +23,9 @@ THE SOFTWARE. import loopy as lp import numpy as np +import sumpy.symbolic as sym from pymbolic.mapper import IdentityMapper, CSECachingMapperMixin from sumpy.symbolic import pymbolic_real_norm_2 -import sumpy.symbolic as sym from pymbolic.primitives import make_sym_vector from pymbolic import var, parse from pytools import memoize_method @@ -211,7 +211,7 @@ class Kernel: expr = expr.diff(vec[i], mi[i]) return expr - def postprocess_at_source(self, expr, avec): + def postprocess_at_source(self, expr, avec, bvec=None): """Transform a kernel evaluation or expansion expression in a place where the vector a (something - source) is known. ("something" may be an expansion center or a target.) @@ -219,12 +219,20 @@ class Kernel: The typical use of this function is to apply source-variable derivatives to the kernel. """ - from sumpy.tools import (ExprDerivativeTaker, - DifferentiatedExprDerivativeTaker) + from sumpy.tools import ( + ExprDerivativeTaker, + DifferentiatedExprDerivativeTaker, + AnisotropicDifferentiatedExprDerivativeTaker) + expr_dict = {(0,)*self.dim: 1} expr_dict = self.get_derivative_transformation_at_source(expr_dict) + if isinstance(expr, ExprDerivativeTaker): - return DifferentiatedExprDerivativeTaker(expr, expr_dict) + if bvec is not None: + # qbmax + return AnisotropicDifferentiatedExprDerivativeTaker(expr) + else: + return DifferentiatedExprDerivativeTaker(expr, expr_dict) result = 0 for mi, coeff in expr_dict.items(): @@ -1264,6 +1272,329 @@ class TargetPointMultiplier(KernelWrapper): # }}} +# {{{ rescaled kernels used by QBMAX + +class _AsymptoticallyRescaledKernelExpressionFactory: + """The rescaled kernels depend on expansion methods. This class isolates + mathematical asymptotic expansions that depend solely on the kernel and + abstract geometry. + + .. attribute :: dim + + The dimension of the ambient space. + + .. attribute :: kernel + + A :class:`~sumpy.kernel.Kernel` object. + + .. attribute :: expr + + The asymptotic expansions. A scalar expression that depends on distance + from the boundary (denoted by the symbol "dist"), satisfying the + normalizing condition: + + expr(dist=0) = 1 + + .. attribute :: geometric_order + + Order of geometric approximation when computing the distance to + boundary. For example, when order=1, the distance is calculated using a + first-order approximation to the boundary using QBX disk's radius and + the point of tangency. + """ + + def __init__(self, kernel, expr, geometric_order=1): + self.dim = kernel.dim + self.kernel = kernel + self.geometric_order = geometric_order + + from sympy import Expr + if isinstance(expr, Expr): + from sumpy.symbolic import SympyToPymbolicMapper + self._expr = SympyToPymbolicMapper()(expr) + else: + from pymbolic.primitives import Expression + assert isinstance(expr, (Expression, int)) + self._expr = expr + + @property + def expr(self): + from sumpy.symbolic import PymbolicToSympyMapperWithSymbols + return PymbolicToSympyMapperWithSymbols()(self._expr) + + @memoize_method + def _get_scaling_at_target(self, expn_class): + """Target-dependent scaling expression, used for scaling back the + potential after summing up the expansion. + """ + full_scaling = self._get_scaling_for_expansion(expn_class) + + from sumpy.expansion.local import ( + LineTaylorLocalExpansion, VolumeTaylorLocalExpansionBase) + if issubclass(expn_class, LineTaylorLocalExpansion): + sym_map = {sym.sym.Symbol("tau"): 1} # target is at tau=1 + elif issubclass(expn_class, VolumeTaylorLocalExpansionBase): + sym_map = {} + else: + raise ValueError("unsupported expansion class") + + return full_scaling.subs(sym_map) + + @memoize_method + def _get_scaling_for_expansion(self, expn_class): + """Target-dependent scaling expression, including the virtual variables + like "tau" for taking expansions. + """ + sym_map = self._get_symbol_map(expn_class) + expr_h = self._approximate_expr() + scaling = expr_h.subs(sym_map) + return scaling + + @memoize_method + def _get_rescaled_kernel_expr(self, expn_class): + """Apply scaling to the kernel expression. + """ + kernel_expr = self._get_kernel_expression_at_source(expn_class) + scaling = self._get_scaling_for_expansion(expn_class) + return kernel_expr / scaling + + @memoize_method + def _approximate_expr(self): + """Transform the expression by approximating the distance with + geometric information known to the QBX routine. + """ + normal, target, dist, rad = ( + self._nvec, self._tvec, self._dist, self._radius) + + if self.geometric_order == 1: + dist_approx = rad - normal.dot(target) + else: + # need curvature information for higher order approximations + raise NotImplementedError + + return self.expr.subs(dist, dist_approx) + + @memoize_method + def _get_symbol_map(self, expn_class): + """Get interpretation of ``nvec`` and ``tvec`` that depends on the + expansion class. + """ + from sumpy.expansion.local import ( + LineTaylorLocalExpansion, VolumeTaylorLocalExpansionBase) + + bvec = sym.make_sym_vector("b", self.dim) + + if issubclass(expn_class, LineTaylorLocalExpansion): + # in a line expansion, bvec is the normal, and the actual target is + # identified with tau. + + bvec_norm = sym.sym.sqrt(bvec.dot(bvec)) + sym_map = { + self._radius: bvec_norm + } + + # self._nvec: bvec / bvec_norm + for lhs, rhs in zip(self._nvec, bvec): + sym_map[lhs] = rhs / bvec_norm + + # self._tvec: sym.Symbol("tau") * bvec + for lhs, rhs in zip(self._tvec, bvec): + sym_map[lhs] = sym.sym.Symbol("tau") * rhs + + elif issubclass(expn_class, VolumeTaylorLocalExpansionBase): + sym_map = { + # disk radius is used as rscale + self._radius: sym.sym.Symbol("rscale"), + } + + # self._nvec: normal / norm(normal) + normal = sym.make_sym_vector("normal", self.dim) + normal_norm = sym.sym.sqrt(normal.dot(normal)) + for lhs, rhs in zip(self._nvec, normal): + sym_map[lhs] = rhs / normal_norm + + # self._tvec: bvec + for lhs, rhs in zip(self._tvec, bvec): + sym_map[lhs] = rhs + + else: + raise ValueError("unsupported expansion class") + + return sym_map + + @memoize_method + def _get_kernel_expression_at_source(self, expn_class): + """Kernel expression used by ``coefficients_from_source``. + """ + from sumpy.expansion.local import ( + LineTaylorLocalExpansion, VolumeTaylorLocalExpansionBase) + + avec = sym.make_sym_vector("a", self.dim) + bvec = sym.make_sym_vector("b", self.dim) + + if issubclass(expn_class, LineTaylorLocalExpansion): + tau = sym.sym.Symbol("tau") + avec_line = avec + tau*bvec + return self.kernel.get_expression(avec_line) + elif issubclass(expn_class, VolumeTaylorLocalExpansionBase): + rad = avec + bvec + return self.kernel.get_expression(rad) + else: + raise ValueError( + f"unsupported expansion class for QBMAX: {expn_class.__name__}") + + @property + def _dist(self): + """The distance of the target point from the boundary. + """ + return sym.sym.Symbol("dist") + + @property + def _nvec(self): + """The unit normal vector obtained by normalizing the vector from the + QBX expansion center to the point of tangency. + """ + return sym.make_sym_vector("nvec", self.dim) + + @property + def _tvec(self): + """The vector from the QBX expansion center to the evaluation target + point. + """ + return sym.make_sym_vector("tvec", self.dim) + + @property + def _radius(self): + """Radius of the QBX disk, aka distance of the expansion center from + the boundary. + """ + return sym.sym.Symbol("radius") + + +class AsymptoticallyInformedKernel(KernelWrapper): + """An asymptotically informed kernel should behave the same as its inner kernel + if the handler does not know how to take advantage of the added information. + Specifically, the kernel keeps a record of what type of expansion will be used + on it. This information needs to be kept updated by the FMM driver for correct + scaling. + """ + init_arg_names = ("inner_kernel", "scaling_expression", "geometric_order", + "expansion_class") + + def __init__(self, inner_kernel, scaling_expression, geometric_order=1, + expansion_class=None): + """ + :param inner_kernel: A raw PDE kernel being rescaled. + :param scaling_expression: A pymbolic/sympy expression. The multiplier + used for scaling the inner_kernel. The scaling is a function + of "dist". + :param geometric_order: An integer for the geometric order. + :param expansion_class: Expansion class to be used with the kernel. + """ + super().__init__(inner_kernel) + + from sympy import Expr + if isinstance(scaling_expression, Expr): + from sumpy.symbolic import SympyToPymbolicMapper + self.scaling_expression = SympyToPymbolicMapper()(scaling_expression) + else: + from pymbolic.primitives import Expression + assert isinstance(scaling_expression, (Expression, int)) + self.scaling_expression = scaling_expression + + self.geometric_order = geometric_order + self.asymptotics = _AsymptoticallyRescaledKernelExpressionFactory( + inner_kernel, scaling_expression, geometric_order) + + self.expansion_class = expansion_class + + def __repr__(self): + return "AsymKnl[%s,%s]" % (self.inner_kernel.__repr__(), + repr(self.expansion_class)) + + def __getinitargs__(self): + return (self.inner_kernel, self.scaling_expression, + self.geometric_order, self.expansion_class) + + def replace_inner_kernel(self, new_inner_kernel): + return type(self)(new_inner_kernel, + self.scaling_expression, self.geometric_order, + self.expansion_class) + + def replace_expansion_class(self, expn_class): + return type(self)(self.inner_kernel, + self.scaling_expression, self.geometric_order, + expn_class) + + def get_base_kernel(self): + """The base kernel is wrapped with the same asymptotic information. + """ + return type(self)( + self.inner_kernel.get_base_kernel(), self.scaling_expression, + self.geometric_order, self.expansion_class) + + def get_expression(self, scaled_dist_vec, expn_class=None): + """To get the asymptotically rescaled expression, pass in + scaled_dist_vec = *None*. Otherwise, it delegates the request to its + inner kernel. + """ + # In these special cases, we allow falling back to the inner kernel + # to be compatible with p2p, p2e and e2p. It is important to let + # other forms like a+tau*b be rescaled. + # In the meanwhile, be careful not to include the scaling into the p2p + # result. + from sumpy.symbolic import make_sym_vector + avec = make_sym_vector("a", self.dim) + bvec = make_sym_vector("b", self.dim) + dvec = make_sym_vector("d", self.dim) + if scaled_dist_vec in [dvec, bvec, avec]: + return self.inner_kernel.get_expression(scaled_dist_vec) + + # Return properly rescaled expression based on expansion class + if not expn_class: + expn_class = self.expansion_class + if not expn_class: + raise ValueError( + "In order to generate expression, the asymptotic expansion " + "needs to know the expansion class.") + + return self.asymptotics._get_rescaled_kernel_expr(expn_class) + + def get_scaling_expression(self, scaled_dist_vec, expn_class=None): + if scaled_dist_vec: + # the asymptotic expression has its own opinion on the dist vec + raise ValueError + if not expn_class: + expn_class = self.expansion_class + if not expn_class: + raise ValueError( + "In order to generate expression, the asymptotic expansion " + "needs to know the expansion class.") + return self.asymptotics._get_scaling_at_target(expn_class) + + def get_pde_as_diff_op(self): + r"""Nonlinear scaling breaks the PDE. + TODO: compression is still possible for the bulk FMM, by using full + Taylor expansion for only qbmax locals. + """ + raise NotImplementedError + + def update_persistent_hash(self, key_hash, key_builder): + for name, value in zip(self.init_arg_names, self.__getinitargs__()): + if name in ["scaling_expression"]: + from pymbolic.mapper.persistent_hash import ( + PersistentHashWalkMapper as PersistentHashWalkMapper) + PersistentHashWalkMapper(key_hash)(value) + elif name in ["expansion_class"]: + key_builder.rec(key_hash, repr(value)) + else: + key_builder.rec(key_hash, value) + + mapper_method = "map_asymptotically_informed_kernel" + +# }}} + + # {{{ kernel mappers class KernelMapper: @@ -1305,6 +1636,7 @@ class KernelIdentityMapper(KernelMapper): map_elasticity_kernel = map_expression_kernel map_line_of_compression_kernel = map_expression_kernel map_stresslet_kernel = map_expression_kernel + map_asymptotically_informed_kernel = map_expression_kernel def map_axis_target_derivative(self, kernel): return type(kernel)(kernel.axis, self.rec(kernel.inner_kernel)) @@ -1319,6 +1651,13 @@ class KernelIdentityMapper(KernelMapper): map_directional_source_derivative = map_directional_target_derivative + def map_asymptotically_informed_kernel(self, kernel): + mapped = type(kernel)( + self.rec(kernel.inner_kernel), + kernel.scaling_expression, kernel.geometric_order, + kernel.expansion_class) + return mapped + class AxisSourceDerivativeRemover(KernelIdentityMapper): def map_axis_source_derivative(self, kernel): @@ -1360,7 +1699,9 @@ class DerivativeCounter(KernelCombineMapper): map_helmholtz_kernel = map_expression_kernel map_yukawa_kernel = map_expression_kernel map_line_of_compression_kernel = map_expression_kernel + map_elasticity_kernel = map_expression_kernel map_stresslet_kernel = map_expression_kernel + map_asymptotically_informed_kernel = map_expression_kernel def map_axis_target_derivative(self, kernel): return 1 + self.rec(kernel.inner_kernel) @@ -1369,6 +1710,47 @@ class DerivativeCounter(KernelCombineMapper): map_directional_source_derivative = map_axis_target_derivative map_axis_source_derivative = map_axis_target_derivative + +class AsymptoticScalingRemover(KernelIdentityMapper): + def map_asymptotically_informed_kernel(self, kernel): + return self.rec(kernel.inner_kernel) + + +class AsymptoticScalingGetter(KernelCombineMapper): + """This is needed because wrappers are typically ordered from + inside out as: + - base kernel + - source derivatives + - asymptotic information + - target derivatives + """ + def combine(self, values): + return values[0] + + def map_expression_kernel(self, kernel): + return None + + map_laplace_kernel = map_expression_kernel + map_biharmonic_kernel = map_expression_kernel + map_helmholtz_kernel = map_expression_kernel + map_yukawa_kernel = map_expression_kernel + map_line_of_compression_kernel = map_expression_kernel + map_elasticity_kernel = map_expression_kernel + map_stresslet_kernel = map_expression_kernel + + def map_asymptotically_informed_kernel(self, kernel): + return kernel.get_scaling_expression(None) + + +class ExpaniosnClassReplacer(KernelIdentityMapper): + def __init__(self, expansion_class): + self.expansion_class = expansion_class + + def map_asymptotically_informed_kernel(self, kernel): + mapped = kernel.replace_expansion_class(self.expansion_class) + mapped_inner = self.rec(kernel.inner_kernel) + return mapped.replace_inner_kernel(mapped_inner) + # }}} diff --git a/sumpy/p2e.py b/sumpy/p2e.py index cb4b6616dca435db2c77350f5aaf5952dacdf6d3..ddb4353abfb7ac8539863a2620364ad32267f081 100644 --- a/sumpy/p2e.py +++ b/sumpy/p2e.py @@ -62,8 +62,9 @@ class P2EBase(KernelComputation, KernelCacheWrapper): number of strength arrays that need to be passed in. By default all kernels use the same strength. """ - from sumpy.kernel import (TargetTransformationRemover, - SourceTransformationRemover) + from sumpy.kernel import ( + TargetTransformationRemover, + SourceTransformationRemover, AsymptoticallyInformedKernel) txr = TargetTransformationRemover() sxr = SourceTransformationRemover() @@ -76,20 +77,43 @@ class P2EBase(KernelComputation, KernelCacheWrapper): for knl in kernels: assert txr(knl) == knl - assert sxr(knl) == expansion.kernel + if sxr(knl) != expansion.kernel: + assert isinstance(sxr(knl), AsymptoticallyInformedKernel) + assert isinstance(expansion.kernel, AsymptoticallyInformedKernel) + assert sxr(knl).inner_kernel == expansion.kernel.inner_kernel + + kernels = [knl.replace_expansion_class(type(expansion)) + if isinstance(knl, AsymptoticallyInformedKernel) + else knl + for knl in kernels] KernelComputation.__init__(self, ctx=ctx, target_kernels=[], source_kernels=kernels, strength_usage=strength_usage, value_dtypes=None, name=name, device=device) + # TODO: target derivatives of scaled kernels using Leibniz's rule + # (fg)' = f'g + fg' ==> f' = [(fg)' - fg'] / g = (fg)'/g - fg'/g, + # i.e., + # f' ~ QBX[(fg)']/g - QBX[fg]*g'/g^2, + # where + # QBX[(fg)'] ~ QBX[fg]', + # can be computed by differentiating QBMAX local expansions. + self.expansion = expansion self.dim = expansion.dim + self._vector_names = {"a"} - def get_loopy_instructions(self): + def get_loopy_instructions(self, isotropic_expansion=True): from sumpy.symbolic import make_sym_vector avec = make_sym_vector("a", self.dim) + if isotropic_expansion: + bvec = None + else: + # non-isotropic expansions like QBMAX depend on bvec + bvec = make_sym_vector("b", self.dim) + import sumpy.symbolic as sp rscale = sp.Symbol("rscale") @@ -97,8 +121,9 @@ class P2EBase(KernelComputation, KernelCacheWrapper): sac = SymbolicAssignmentCollection() strengths = [sp.Symbol(f"strength_{i}") for i in self.strength_usage] + coeffs = self.expansion.coefficients_from_source_vec(self.source_kernels, - avec, None, rscale, strengths, sac=sac) + avec, bvec, rscale, strengths, sac=sac) coeff_names = [] for i, coeff in enumerate(coeffs): @@ -113,7 +138,7 @@ class P2EBase(KernelComputation, KernelCacheWrapper): from sumpy.codegen import to_loopy_insns return to_loopy_insns( sac.assignments.items(), - vector_names={"a"}, + vector_names=self._vector_names, pymbolic_expr_maps=code_transformers, retain_names=coeff_names, ) @@ -369,5 +394,3 @@ class P2EFromCSR(P2EBase): return super().__call__(queue, **kwargs) # }}} - -# vim: foldmethod=marker diff --git a/sumpy/p2p.py b/sumpy/p2p.py index c2aa895d5dec0c250149fe5166ede72cf8534e68..f7625a401aa42c087b61bcf63389ef8fc50f7ae5 100644 --- a/sumpy/p2p.py +++ b/sumpy/p2p.py @@ -88,6 +88,7 @@ class P2PBase(KernelComputation, KernelCacheWrapper): self.exclude_self = exclude_self + self._vector_names = {"d"} self.dim = single_valued(knl.dim for knl in list(self.target_kernels) + list(self.source_kernels)) @@ -134,7 +135,7 @@ class P2PBase(KernelComputation, KernelCacheWrapper): from sumpy.codegen import to_loopy_insns loopy_insns = to_loopy_insns(sac.assignments.items(), - vector_names={"d"}, + vector_names=self._vector_names, pymbolic_expr_maps=( [knl.get_code_transformer() for knl in self.source_kernels] + [knl.get_code_transformer() for knl in self.target_kernels]), diff --git a/sumpy/qbx.py b/sumpy/qbx.py index 30d3254f1a65f93074fdfd586ade54b3b9907b37..be9d55a6d083c16e103b55f1f5ea0b36e6afc529 100644 --- a/sumpy/qbx.py +++ b/sumpy/qbx.py @@ -69,16 +69,22 @@ def stringify_expn_index(i): class LayerPotentialBase(KernelComputation, KernelCacheWrapper): def __init__(self, ctx, expansion, strength_usage=None, value_dtypes=None, name=None, device=None, - source_kernels=None, target_kernels=None): + source_kernels=None, target_kernels=None, _use_qbmax=False): - from pytools import single_valued + from sumpy.kernel import ExpaniosnClassReplacer + ecr = ExpaniosnClassReplacer(type(expansion)) + source_kernels = tuple(ecr(sk) for sk in source_kernels) + target_kernels = tuple(ecr(tk) for tk in target_kernels) + from pytools import single_valued KernelComputation.__init__(self, ctx=ctx, target_kernels=target_kernels, strength_usage=strength_usage, source_kernels=source_kernels, value_dtypes=value_dtypes, name=name, device=device) self.dim = single_valued(knl.dim for knl in self.target_kernels) self.expansion = expansion + self._vector_names = {"a", "b"} + self._use_qbmax = _use_qbmax def get_cache_key(self): return (type(self).__name__, self.expansion, tuple(self.target_kernels), @@ -115,8 +121,14 @@ class LayerPotentialBase(KernelComputation, KernelCacheWrapper): coefficients[self.expansion.get_storage_index(i)])) for i in self.expansion.get_coefficient_identifiers()] + if self._use_qbmax is None or self._use_qbmax: + use_qbmax = True + else: + use_qbmax = False + return sac.assign_unique("expn%d_result" % expansion_nr, - self.expansion.evaluate(tgt_knl, assigned_coeffs, bvec, rscale)) + self.expansion.evaluate(tgt_knl, assigned_coeffs, bvec, rscale, + use_qbmax=use_qbmax)) def get_loopy_insns_and_result_names(self): from sumpy.symbolic import make_sym_vector @@ -145,7 +157,7 @@ class LayerPotentialBase(KernelComputation, KernelCacheWrapper): from sumpy.codegen import to_loopy_insns loopy_insns = to_loopy_insns( sac.assignments.items(), - vector_names={"a", "b"}, + vector_names=self._vector_names, pymbolic_expr_maps=pymbolic_expr_maps, retain_names=result_names, complex_dtype=np.complex128 # FIXME @@ -229,8 +241,6 @@ class LayerPotential(LayerPotentialBase): @memoize_method def get_kernel(self): - loopy_insns, result_names = self.get_loopy_insns_and_result_names() - kernel_exprs = self.get_kernel_exprs(result_names) arguments = ( self.get_default_src_tgt_arguments() + [lp.GlobalArg("strength_%d" % i, @@ -240,6 +250,56 @@ class LayerPotential(LayerPotentialBase): None, shape="ntargets", order="C") for i in range(len(self.target_kernels))]) + sats = [] + sat_insns_data = [] + + for iknl, knl in enumerate(self.target_kernels): + from sumpy.kernel import AsymptoticScalingGetter, ExpaniosnClassReplacer + ecr = ExpaniosnClassReplacer(type(self.expansion)) + knl = ecr(knl) + + asg = AsymptoticScalingGetter() + sat_expr = asg(knl) + + if sat_expr: + # inject the multiplicative scales of each output kernel + sats.append(f"sat_{iknl} * ") + sat_insns_data.append((f"sat_{iknl}", sat_expr)) + else: + sats.append("") + + qbmax_local_temp_insns = [] + if sat_insns_data: + # Volume expansion needs normal vector information. + # Though using volume expansion, it only evaluates self interaction. + # Centers and targets are 1-to-1. + from sumpy.expansion import VolumeTaylorExpansionBase + if isinstance(self.expansion, VolumeTaylorExpansionBase): + qbmax_local_temp_insns = [""" + <> normal[idim] = \ + nodes_of_qbx_balls_tangency[idim, itgt//2] - \ + center[idim, itgt] {dup=idim} + """] + + arguments += ( + lp.GlobalArg("nodes_of_qbx_balls_tangency", None, + shape=(self.dim, "ntargets"), dim_tags="sep,C"), + ) + + self._vector_names.add("normal") + + from sumpy.codegen import to_loopy_insns + sat_insns = to_loopy_insns( + sat_insns_data, + vector_names=self._vector_names, + retain_names=[datpair[0] for datpair in sat_insns_data], + complex_dtype=np.complex128) # FIXME + else: + sat_insns = [] + + loopy_insns, result_names = self.get_loopy_insns_and_result_names() + kernel_exprs = self.get_kernel_exprs(result_names) + loopy_knl = lp.make_kernel([""" {[itgt, isrc, idim]: \ 0 <= itgt < ntargets and \ @@ -253,12 +313,13 @@ class LayerPotential(LayerPotentialBase): + ["<> rscale = expansion_radii[itgt]"] + [f"<> strength_{i}_isrc = strength_{i}[isrc]" for i in range(self.strength_count)] + + sat_insns + qbmax_local_temp_insns + loopy_insns + kernel_exprs + [""" - result_{i}[itgt] = knl_{i}_scaling * \ + result_{i}[itgt] = knl_{i}_scaling * {sat} \ simul_reduce(sum, isrc, pair_result_{i}) \ {{id_prefix=write_lpot,inames=itgt}} - """.format(i=iknl) + """.format(i=iknl, sat=sats[iknl]) for iknl in range(len(self.target_kernels))] + ["end"], arguments, @@ -281,7 +342,6 @@ class LayerPotential(LayerPotentialBase): :arg strengths: are required to have area elements and quadrature weights already multiplied in. """ - knl = self.get_cached_optimized_kernel( targets_is_obj_array=is_obj_array_like(targets), sources_is_obj_array=is_obj_array_like(sources), diff --git a/sumpy/tools.py b/sumpy/tools.py index 5b3e31308e72e28a8e6491d6cb83d1a3824b71be..4ecb048d39077eb3564ed409fca877ca190ca41a 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -412,7 +412,7 @@ class DifferentiatedExprDerivativeTaker: taker: ExprDerivativeTaker derivative_transformation: dict - def diff(self, mi, save_intermediate=lambda x: x): + def diff(self, mi, save_intermediate=lambda x: x, subs=None): # By passing `rscale` to the derivative taker we are taking a scaled # version of the derivative which is `expr.diff(mi)*rscale**sum(mi)` # which might be implemented efficiently for kernels like Laplace. @@ -428,8 +428,34 @@ class DifferentiatedExprDerivativeTaker: / self.taker.rscale ** (sum(extra_mi) - max_order) for extra_mi, coeff in self.derivative_transformation.items()) + if subs: + result = result.subs(subs) + return result * save_intermediate(1 / self.taker.rscale ** max_order) + +@tag_dataclass +class AnisotropicDifferentiatedExprDerivativeTaker: + """Implements the :class:`ExprDerivativeTaker` interface for an expression + that is itself a linear combination of derivatives of a base expression. To + take the actual derivatives, it makes use of an underlying derivative taker + *taker*. + + Unlike :class:`DifferentiatedExprDerivativeTaker`, the derivatives inside + the taker are already evaluated. Because for anisotropic (target-dependent) + qbmax expansions, derivative transforms are nonlinear. + + .. attribute:: taker + A :class:`ExprDerivativeTaker` for the base expression. + """ + taker: ExprDerivativeTaker + + def diff(self, mi, save_intermediate=lambda x: x, subs=None): + result = self.taker.diff(mi) + if subs: + result = result.subs(subs) + return result * save_intermediate(1) + # }}} diff --git a/test/test_qbx.py b/test/test_qbx.py index 379ef0e221327f13d5452c34e8f1601de9fad421..d23f74d759c857d878bfb6a5697fb364f197345f 100644 --- a/test/test_qbx.py +++ b/test/test_qbx.py @@ -33,12 +33,22 @@ from sumpy.expansion.local import ( LineTaylorLocalExpansion, VolumeTaylorLocalExpansion) import pytest +from sumpy.kernel import LaplaceKernel, AsymptoticallyInformedKernel -@pytest.mark.parametrize("expn_class", [ - LineTaylorLocalExpansion, - VolumeTaylorLocalExpansion, +import sympy as sp +dist = sp.Symbol("dist") +asexpr = sp.exp(-dist) # no practical use, just to test the machinary + + +@pytest.mark.parametrize("expn_class,lknl", [ + (LineTaylorLocalExpansion, LaplaceKernel(2)), + (VolumeTaylorLocalExpansion, LaplaceKernel(2)), + (LineTaylorLocalExpansion, AsymptoticallyInformedKernel( + LaplaceKernel(2), asexpr)), + (VolumeTaylorLocalExpansion, AsymptoticallyInformedKernel( + LaplaceKernel(2), asexpr)), ]) -def test_direct_qbx_vs_eigval(ctx_factory, expn_class): +def test_direct_qbx_vs_eigval(ctx_factory, expn_class, lknl): """This evaluates a single layer potential on a circle using a known eigenvalue/eigenvector combination. """ @@ -48,15 +58,15 @@ def test_direct_qbx_vs_eigval(ctx_factory, expn_class): ctx = ctx_factory() queue = cl.CommandQueue(ctx) - from sumpy.kernel import LaplaceKernel - lknl = LaplaceKernel(2) - order = 12 from sumpy.qbx import LayerPotential - lpot = LayerPotential(ctx, expansion=expn_class(lknl, order), - target_kernels=(lknl,), source_kernels=(lknl,)) + use_qbmax = isinstance(lknl, AsymptoticallyInformedKernel) + lpot = LayerPotential( + ctx, expansion=expn_class(lknl, order), + target_kernels=(lknl,), source_kernels=(lknl,), + _use_qbmax=use_qbmax) mode_nr = 25 @@ -84,9 +94,16 @@ def test_direct_qbx_vs_eigval(ctx_factory, expn_class): expansion_radii = np.ones(n) * radius + if use_qbmax and (expn_class is VolumeTaylorLocalExpansion): + qbmax_extra_kwargs = {"nodes_of_qbx_balls_tangency": unit_circle} + else: + qbmax_extra_kwargs = dict() + strengths = (sigma * h,) - evt, (result_qbx,) = lpot(queue, targets, sources, centers, strengths, - expansion_radii=expansion_radii) + evt, (result_qbx,) = lpot( + queue, targets, sources, centers, strengths, + expansion_radii=expansion_radii, + **qbmax_extra_kwargs) eocrec.add_data_point(h, np.max(np.abs(result_ref - result_qbx))) @@ -96,11 +113,15 @@ def test_direct_qbx_vs_eigval(ctx_factory, expn_class): assert eocrec.order_estimate() > order - slack -@pytest.mark.parametrize("expn_class", [ - LineTaylorLocalExpansion, - VolumeTaylorLocalExpansion, +@pytest.mark.parametrize("expn_class, lknl", [ + (LineTaylorLocalExpansion, LaplaceKernel(2)), + (VolumeTaylorLocalExpansion, LaplaceKernel(2)), + (LineTaylorLocalExpansion, AsymptoticallyInformedKernel( + LaplaceKernel(2), asexpr)), + (VolumeTaylorLocalExpansion, AsymptoticallyInformedKernel( + LaplaceKernel(2), asexpr)), ]) -def test_direct_qbx_vs_eigval_with_tgt_deriv(ctx_factory, expn_class): +def test_direct_qbx_vs_eigval_with_tgt_deriv(ctx_factory, expn_class, lknl): """This evaluates a single layer potential on a circle using a known eigenvalue/eigenvector combination. """ @@ -110,17 +131,22 @@ def test_direct_qbx_vs_eigval_with_tgt_deriv(ctx_factory, expn_class): ctx = ctx_factory() queue = cl.CommandQueue(ctx) - from sumpy.kernel import LaplaceKernel, AxisTargetDerivative - lknl = LaplaceKernel(2) + from sumpy.kernel import AxisTargetDerivative order = 8 from sumpy.qbx import LayerPotential - lpot_dx = LayerPotential(ctx, expansion=expn_class(lknl, order), - target_kernels=(AxisTargetDerivative(0, lknl),), source_kernels=(lknl,)) - lpot_dy = LayerPotential(ctx, expansion=expn_class(lknl, order), - target_kernels=(AxisTargetDerivative(1, lknl),), source_kernels=(lknl,)) + # explicitly enable _use_qbmax when directly using LayerPotential + use_qbmax = isinstance(lknl, AsymptoticallyInformedKernel) + lpot_dx = LayerPotential( + ctx, expansion=expn_class(lknl, order), + target_kernels=(AxisTargetDerivative(0, lknl),), source_kernels=(lknl,), + _use_qbmax=use_qbmax) + lpot_dy = LayerPotential( + ctx, expansion=expn_class(lknl, order), + target_kernels=(AxisTargetDerivative(1, lknl),), source_kernels=(lknl,), + _use_qbmax=use_qbmax) mode_nr = 15 @@ -151,10 +177,15 @@ def test_direct_qbx_vs_eigval_with_tgt_deriv(ctx_factory, expn_class): strengths = (sigma * h,) + if use_qbmax and (expn_class is VolumeTaylorLocalExpansion): + qbmax_extra_kwargs = {"nodes_of_qbx_balls_tangency": unit_circle} + else: + qbmax_extra_kwargs = dict() + evt, (result_qbx_dx,) = lpot_dx(queue, targets, sources, centers, strengths, - expansion_radii=expansion_radii) + expansion_radii=expansion_radii, **qbmax_extra_kwargs) evt, (result_qbx_dy,) = lpot_dy(queue, targets, sources, centers, strengths, - expansion_radii=expansion_radii) + expansion_radii=expansion_radii, **qbmax_extra_kwargs) normals = unit_circle result_qbx = normals[0] * result_qbx_dx + normals[1] * result_qbx_dy