diff --git a/sumpy/assignment_collection.py b/sumpy/assignment_collection.py index 840b04da0c65ad4a425cee47ab88964caf8a424c..8aa07d813066f6dc0382a57e34c42f14ae8bced7 100644 --- a/sumpy/assignment_collection.py +++ b/sumpy/assignment_collection.py @@ -26,7 +26,7 @@ THE SOFTWARE. """ -import sympy as sp +import sumpy.symbolic as sym import logging logger = logging.getLogger(__name__) @@ -63,7 +63,7 @@ class _SymbolGenerator(object): self.base_to_count[base] = count + 1 - return sp.Symbol(id_str) + return sym.Symbol(id_str) def __iter__(self): return self @@ -125,7 +125,7 @@ class SymbolicAssignmentCollection(object): result = set() for dep in self.assignments[var_name].atoms(): - if not isinstance(dep, sp.Symbol): + if not isinstance(dep, sym.Symbol): continue dep_name = dep.name @@ -146,7 +146,7 @@ class SymbolicAssignmentCollection(object): if root_name is None: root_name = name - self.assignments[name] = sp.sympify(expr) + self.assignments[name] = sym.sympify(expr) def assign_unique(self, name_base, expr): """Assign *expr* to a new variable whose name is based on *name_base*. @@ -170,7 +170,7 @@ class SymbolicAssignmentCollection(object): # Options here: # - checked_cse: if you mistrust the result of the cse. # Uses maxima to verify. - # - sp.cse: The sympy thing. + # - sym.cse: The sympy thing. # - sumpy.cse.cse: Based on sympy, designed to go faster. #from sumpy.symbolic import checked_cse @@ -185,41 +185,13 @@ class SymbolicAssignmentCollection(object): self.assignments[name] = new_expr for name, value in new_assignments: - assert isinstance(name, sp.Symbol) + assert isinstance(name, sym.Symbol) self.add_assignment(name.name, value) logger.info("common subexpression elimination: done after {dur:.2f} s" .format(dur=time.time() - start_time)) return new_extra_exprs - def kill_trivial_assignments(self, exprs): - logger.info("kill trivial assignments: start") - - approved_assignments = [] - rejected_assignments = [] - - from sumpy.symbolic import is_assignment_nontrivial - for name, value in six.iteritems(self.assignments): - if name in self.user_symbols or is_assignment_nontrivial(name, value): - approved_assignments.append((name, value)) - else: - rejected_assignments.append((name, value)) - - # un-substitute rejected assignments - from sumpy.symbolic import make_one_step_subst - unsubst_rej = make_one_step_subst(rejected_assignments) - - new_assignments = dict( - (name, expr.subs(unsubst_rej)) - for name, expr in approved_assignments) - - exprs = [expr.subs(unsubst_rej) for expr in exprs] - self.assignments = new_assignments - - logger.info("kill trivial assignments: done") - - return exprs - # }}} # vim: fdm=marker diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 0cb14028261ca7d8a977c26f340c5606f6fd26e3..59cc9071ed17c8d0b56ffef74bfd1cfde4ae5371 100644 --- a/sumpy/codegen.py +++ b/sumpy/codegen.py @@ -38,7 +38,7 @@ from loopy.types import NumpyType from pytools import memoize_method -from pymbolic.interop.sympy import ( +from sumpy.symbolic import (USE_SYMENGINE, SympyToPymbolicMapper as SympyToPymbolicMapperBase) import logging @@ -58,6 +58,10 @@ Conversion of :mod:`sympy` expressions to :mod:`loopy` # {{{ sympy -> pymbolic mapper +import sumpy.symbolic as sym +_SPECIAL_FUNCTION_NAMES = frozenset(dir(sym.functions)) + + class SympyToPymbolicMapper(SympyToPymbolicMapperBase): def __init__(self, assignments): self.assignments = assignments @@ -77,6 +81,24 @@ class SympyToPymbolicMapper(SympyToPymbolicMapperBase): expr.expr.subs(self.assignments)), tuple(v.name for v in expr.variables)) + def not_supported(self, expr): + if isinstance(expr, int): + return expr + elif getattr(expr, "is_Function", False): + if USE_SYMENGINE: + try: + func_name = expr.get_name() + except AttributeError: + func_name = type(expr).__name__ + else: + func_name = type(expr).__name__ + # SymEngine capitalizes the names of the special functions. + if func_name.lower() in _SPECIAL_FUNCTION_NAMES: + func_name = func_name.lower() + return prim.Variable(func_name)( + *tuple(self.rec(arg) for arg in expr.args)) + else: + return SympyToPymbolicMapperBase.not_supported(self, expr) # }}} @@ -318,7 +340,7 @@ class BesselDerivativeReplacer(CSECachingMapperMixin, IdentityMapper): arg, = expr.values n_derivs = len(expr.child.variables) - import sympy as sp + import sympy as sym # AS (9.1.31) # http://dlmf.nist.gov/10.6.7 @@ -329,7 +351,7 @@ class BesselDerivativeReplacer(CSECachingMapperMixin, IdentityMapper): k = n_derivs return prim.CommonSubexpression( 2**(-k)*sum( - (-1)**idx*int(sp.binomial(k, idx)) * function(i, arg) + (-1)**idx*int(sym.binomial(k, idx)) * function(i, arg) for idx, i in enumerate(range(order-k, order+k+1, 2))), "d%d_%s_%s" % (n_derivs, function.name, order_str)) else: diff --git a/sumpy/cse.py b/sumpy/cse.py index 350a2c8d97d6ef4729657c85248d412e4128f6a1..def1c7c6ecebc13bf342ca77d497b8136eaa37a5 100644 --- a/sumpy/cse.py +++ b/sumpy/cse.py @@ -66,8 +66,7 @@ DAMAGE. # }}} -from sympy.core import Basic, Mul, Add, Pow -from sympy.core.function import _coeff_isneg +from sumpy.symbolic import Basic, Mul, Add, Pow, Symbol, _coeff_isneg from sympy.core.compatibility import iterable from sympy.utilities.iterables import numbered_symbols @@ -192,19 +191,16 @@ class FuncArgTracker(object): (k, v) for k, v in count_map.items() if v >= threshold) - def get_subset_candidates(self, argset, restrict_to_funcset=None): + def get_subset_candidates(self, argset, restrict_to_funcset=frozenset()): """ Return a set of functions each of which whose argument list contains - `argset`, optionally filtered only to contain functions in - `restrict_to_funcset`. + `argset.` """ iarg = iter(argset) indices = set( - fi for fi in self.arg_to_funcset[next(iarg)]) - - if restrict_to_funcset is not None: - indices &= restrict_to_funcset + fi for fi in self.arg_to_funcset[next(iarg)] + if fi in restrict_to_funcset) for arg in iarg: indices &= self.arg_to_funcset[arg] @@ -227,6 +223,19 @@ class FuncArgTracker(object): self.func_to_argset[func_i].update(new_args) +class Unevaluated(object): + + def __init__(self, func, args): + self.func = func + self.args = args + + def __str__(self): + return "Uneval<{}>({})".format( + self.func, ", ".join(str(a) for a in self.args)) + + __repr__ = __str__ + + def match_common_args(func_class, funcs, opt_subs): """ Recognize and extract common subexpressions of function arguments within a @@ -284,9 +293,8 @@ def match_common_args(func_class, funcs, opt_subs): diff_i = arg_tracker.func_to_argset[i].difference(com_args) if diff_i: # com_func needs to be unevaluated to allow for recursive matches. - com_func = func_class( - *arg_tracker.get_args_in_value_order(com_args), - evaluate=False) + com_func = Unevaluated( + func_class, arg_tracker.get_args_in_value_order(com_args)) com_func_number = arg_tracker.get_or_add_value_number(com_func) arg_tracker.update_func_argset(i, diff_i | set([com_func_number])) changed.add(i) @@ -313,9 +321,8 @@ def match_common_args(func_class, funcs, opt_subs): changed.add(k) if i in changed: - opt_subs[funcs[i]] = func_class( - *arg_tracker.get_args_in_value_order(arg_tracker.func_to_argset[i]), - evaluate=False) + opt_subs[funcs[i]] = Unevaluated(func_class, + arg_tracker.get_args_in_value_order(arg_tracker.func_to_argset[i])) arg_tracker.stop_arg_tracking(i) @@ -361,7 +368,7 @@ def opt_cse(exprs): if _coeff_isneg(expr): neg_expr = -expr if not neg_expr.is_Atom: - opt_subs[expr] = Mul(-1, neg_expr, evaluate=False) + opt_subs[expr] = Unevaluated(Mul, (-1, neg_expr)) seen_subexp.add(neg_expr) expr = neg_expr @@ -374,7 +381,7 @@ def opt_cse(exprs): elif isinstance(expr, Pow): base, exp = expr.args if _coeff_isneg(exp): - opt_subs[expr] = Pow(Pow(base, -exp), -1, evaluate=False) + opt_subs[expr] = Unevaluated(Pow, (Pow(base, -exp), -1)) # }}} @@ -415,10 +422,10 @@ def tree_cse(exprs, symbols, opt_subs=None): excluded_symbols = set() def find_repeated(expr): - if not isinstance(expr, Basic): + if not isinstance(expr, (Basic, Unevaluated)): return - if expr.is_Atom: + if isinstance(expr, Basic) and expr.is_Atom: if expr.is_Symbol: excluded_symbols.add(expr) return @@ -457,7 +464,7 @@ def tree_cse(exprs, symbols, opt_subs=None): subs = dict() def rebuild(expr): - if not isinstance(expr, Basic): + if not isinstance(expr, (Basic, Unevaluated)): return expr if not expr.args: @@ -475,7 +482,7 @@ def tree_cse(exprs, symbols, opt_subs=None): expr = opt_subs[expr] new_args = [rebuild(arg) for arg in expr.args] - if new_args != expr.args: + if isinstance(expr, Unevaluated) or new_args != expr.args: new_expr = expr.func(*new_args) else: new_expr = expr @@ -539,7 +546,7 @@ def cse(exprs, symbols=None, optimizations=None): reduced_exprs = [preprocess_for_cse(e, optimizations) for e in exprs] if symbols is None: - symbols = numbered_symbols() + symbols = numbered_symbols(cls=Symbol) else: # In case we get passed an iterable with an __iter__ method instead of # an actual iterator. diff --git a/sumpy/e2e.py b/sumpy/e2e.py index f5b349401fee21112ae5e72b28d1d1ec8183046a..f93a868234e0a9468bf865e7eca0c8db2f6fe474 100644 --- a/sumpy/e2e.py +++ b/sumpy/e2e.py @@ -27,7 +27,7 @@ from six.moves import range import numpy as np import loopy as lp -import sympy as sp +import sumpy.symbolic as sym from sumpy.tools import KernelCacheWrapper import logging @@ -77,10 +77,10 @@ class E2EBase(KernelCacheWrapper): self.dim = src_expansion.dim def get_translation_loopy_insns(self): - from sumpy.symbolic import make_sympy_vector - dvec = make_sympy_vector("d", self.dim) + from sumpy.symbolic import make_sym_vector + dvec = make_sym_vector("d", self.dim) - src_coeff_exprs = [sp.Symbol("src_coeff%d" % i) + src_coeff_exprs = [sym.Symbol("src_coeff%d" % i) for i in range(len(self.src_expansion))] from sumpy.assignment_collection import SymbolicAssignmentCollection @@ -202,6 +202,8 @@ class E2EFromCSR(E2EBase): silenced_warnings="write_race(write_expn*)", default_offset=lp.auto) + print(loopy_knl) + loopy_knl = lp.fix_parameters(loopy_knl, dim=self.dim) for expn in [self.src_expansion, self.tgt_expansion]: diff --git a/sumpy/e2p.py b/sumpy/e2p.py index b2e0e58569ad91cc5416f529bc766b9ee7950859..118f1743ca81b5ba02c5f0bc6639e45d91db5eb1 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -27,7 +27,7 @@ from six.moves import range import numpy as np import loopy as lp -import sympy as sp +import sumpy.symbolic as sym from sumpy.tools import KernelCacheWrapper @@ -74,13 +74,13 @@ class E2PBase(KernelCacheWrapper): assert tdr(knl) == expansion.kernel def get_loopy_insns_and_result_names(self): - from sumpy.symbolic import make_sympy_vector - bvec = make_sympy_vector("b", self.dim) + from sumpy.symbolic import make_sym_vector + bvec = make_sym_vector("b", self.dim) from sumpy.assignment_collection import SymbolicAssignmentCollection sac = SymbolicAssignmentCollection() - coeff_exprs = [sp.Symbol("coeff%d" % i) + coeff_exprs = [sym.Symbol("coeff%d" % i) for i in range(len(self.expansion.get_coefficient_identifiers()))] value = self.expansion.evaluate(coeff_exprs, bvec) result_names = [ @@ -107,7 +107,7 @@ class E2PBase(KernelCacheWrapper): return loopy_insns, result_names def get_kernel_scaling_assignment(self): - from pymbolic.interop.sympy import SympyToPymbolicMapper + from sumpy.symbolic import SympyToPymbolicMapper sympy_conv = SympyToPymbolicMapper() return [lp.Assignment(id=None, assignee="kernel_scaling", diff --git a/sumpy/expansion/__init__.py b/sumpy/expansion/__init__.py index cc49c6ca2e712ecbf43a0a6d460d0e0dd2667152..f50bd5fbb9facd732c0c192c787a0a420d72b58c 100644 --- a/sumpy/expansion/__init__.py +++ b/sumpy/expansion/__init__.py @@ -25,9 +25,9 @@ THE SOFTWARE. """ import numpy as np -import sympy as sp import logging from pytools import memoize_method +import sumpy.symbolic as sym from sumpy.tools import MiDerivativeTaker __doc__ = """ @@ -356,7 +356,7 @@ class HelmholtzDerivativeWrangler(LinearRecurrenceBasedDerivativeWrangler): coeffs[needed_deriv] = -1 else: - k = sp.Symbol(self.helmholtz_k_name) + k = sym.Symbol(self.helmholtz_k_name) coeffs[tuple(reduced_deriv)] = -k*k return coeffs diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py index 3ba643cd2e1bdf6af214f069b278f01d650b26d3..e958c2cb38b71d7066f45869a3954a1fc9dcdd37 100644 --- a/sumpy/expansion/local.py +++ b/sumpy/expansion/local.py @@ -23,7 +23,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -import sympy as sp +import sumpy.symbolic as sym from sumpy.expansion import ( ExpansionBase, VolumeTaylorExpansion, LaplaceConformingVolumeTaylorExpansion, @@ -61,7 +61,7 @@ class LineTaylorLocalExpansion(LocalExpansionBase): raise RuntimeError("cannot use line-Taylor expansions in a setting " "where the center-target vector is not known at coefficient " "formation") - avec_line = avec + sp.Symbol("tau")*bvec + avec_line = avec + sym.Symbol("tau")*bvec line_kernel = self.kernel.get_expression(avec_line) @@ -117,8 +117,9 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): type(self).__name__, self.order)) + import time from sumpy.expansion.multipole import VolumeTaylorMultipoleExpansionBase - if isinstance(src_expansion, VolumeTaylorMultipoleExpansionBase): + if 1 and isinstance(src_expansion, VolumeTaylorMultipoleExpansionBase): # We know the general form of the multipole expansion is: # # coeff0 * diff(kernel, mi0) + coeff1 * diff(kernel, mi1) + ... @@ -140,7 +141,7 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): src_expansion.get_coefficient_identifiers()): kernel_deriv = taker.diff(add_mi(deriv, term)) local_result.append(coeff * kernel_deriv) - result.append(sp.Add(*local_result)) + result.append(sym.Add(*local_result)) else: from sumpy.tools import MiDerivativeTaker expr = src_expansion.evaluate(src_coeff_exprs, dvec) @@ -197,61 +198,61 @@ class H2DLocalExpansion(LocalExpansionBase): return list(range(-self.order, self.order+1)) def coefficients_from_source(self, avec, bvec): - from sumpy.symbolic import sympy_real_norm_2 - hankel_1 = sp.Function("hankel_1") + from sumpy.symbolic import sym_real_norm_2 + hankel_1 = sym.Function("hankel_1") - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + k = sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) # The coordinates are negated since avec points from source to center. - source_angle_rel_center = sp.atan2(-avec[1], -avec[0]) - avec_len = sympy_real_norm_2(avec) + source_angle_rel_center = sym.atan2(-avec[1], -avec[0]) + avec_len = sym_real_norm_2(avec) return [self.kernel.postprocess_at_source( hankel_1(l, k * avec_len) - * sp.exp(sp.I * l * source_angle_rel_center), avec) + * sym.exp(sym.I * l * source_angle_rel_center), avec) for l in self.get_coefficient_identifiers()] def evaluate(self, coeffs, bvec): - from sumpy.symbolic import sympy_real_norm_2 - bessel_j = sp.Function("bessel_j") - bvec_len = sympy_real_norm_2(bvec) - target_angle_rel_center = sp.atan2(bvec[1], bvec[0]) + from sumpy.symbolic import sym_real_norm_2 + bessel_j = sym.Function("bessel_j") + bvec_len = sym_real_norm_2(bvec) + target_angle_rel_center = sym.atan2(bvec[1], bvec[0]) - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + k = sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) return sum(coeffs[self.get_storage_index(l)] * self.kernel.postprocess_at_target( bessel_j(l, k * bvec_len) - * sp.exp(sp.I * l * -target_angle_rel_center), bvec) + * sym.exp(sym.I * l * -target_angle_rel_center), bvec) for l in self.get_coefficient_identifiers()) def translate_from(self, src_expansion, src_coeff_exprs, dvec): - from sumpy.symbolic import sympy_real_norm_2 + from sumpy.symbolic import sym_real_norm_2 - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + k = sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) if isinstance(src_expansion, H2DLocalExpansion): - dvec_len = sympy_real_norm_2(dvec) - bessel_j = sp.Function("bessel_j") - new_center_angle_rel_old_center = sp.atan2(dvec[1], dvec[0]) + dvec_len = sym_real_norm_2(dvec) + bessel_j = sym.Function("bessel_j") + new_center_angle_rel_old_center = sym.atan2(dvec[1], dvec[0]) translated_coeffs = [] for l in self.get_coefficient_identifiers(): translated_coeffs.append( sum(src_coeff_exprs[src_expansion.get_storage_index(m)] * bessel_j(m - l, k * dvec_len) - * sp.exp(sp.I * (m - l) * -new_center_angle_rel_old_center) + * sym.exp(sym.I * (m - l) * -new_center_angle_rel_old_center) for m in src_expansion.get_coefficient_identifiers())) return translated_coeffs from sumpy.expansion.multipole import H2DMultipoleExpansion if isinstance(src_expansion, H2DMultipoleExpansion): - dvec_len = sympy_real_norm_2(dvec) - hankel_1 = sp.Function("hankel_1") - new_center_angle_rel_old_center = sp.atan2(dvec[1], dvec[0]) + dvec_len = sym_real_norm_2(dvec) + hankel_1 = sym.Function("hankel_1") + new_center_angle_rel_old_center = sym.atan2(dvec[1], dvec[0]) translated_coeffs = [] for l in self.get_coefficient_identifiers(): translated_coeffs.append( sum((-1) ** l * hankel_1(m + l, k * dvec_len) - * sp.exp(sp.I * (m + l) * new_center_angle_rel_old_center) + * sym.exp(sym.I * (m + l) * new_center_angle_rel_old_center) * src_coeff_exprs[src_expansion.get_storage_index(m)] for m in src_expansion.get_coefficient_identifiers())) return translated_coeffs diff --git a/sumpy/expansion/multipole.py b/sumpy/expansion/multipole.py index fa227e89ab3a0290dd64d9420340ceeec8d93843..ce3b78d8ae25049f48f1b1569ac7483bf8586923 100644 --- a/sumpy/expansion/multipole.py +++ b/sumpy/expansion/multipole.py @@ -23,7 +23,7 @@ THE SOFTWARE. """ from six.moves import range, zip -import sympy as sp # noqa +import sumpy.symbolic as sym # noqa from sumpy.expansion import ( ExpansionBase, VolumeTaylorExpansion, LaplaceConformingVolumeTaylorExpansion, @@ -63,8 +63,8 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): raise NotImplementedError("more than one source derivative " "not supported at present") - from sumpy.symbolic import make_sympy_vector - dir_vec = make_sympy_vector(kernel.dir_vec_name, kernel.dim) + from sumpy.symbolic import make_sym_vector + dir_vec = make_sym_vector(kernel.dir_vec_name, kernel.dim) coeff_identifiers = self.get_full_coefficient_identifiers() result = [0] * len(coeff_identifiers) @@ -91,9 +91,9 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): def evaluate(self, coeffs, bvec): taker = self.get_kernel_derivative_taker(bvec) - result = sum( + result = sym.Add(*tuple( coeff * taker.diff(mi) - for coeff, mi in zip(coeffs, self.get_coefficient_identifiers())) + for coeff, mi in zip(coeffs, self.get_coefficient_identifiers()))) return result def get_kernel_derivative_taker(self, bvec): @@ -143,7 +143,8 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): n = tgt_mi[idim] k = src_mi[idim] assert n >= k - contrib *= (sp.binomial(n, k) + from sympy import binomial + contrib *= (binomial(n, k) * dvec[idim]**(n-k)) result[i] += contrib @@ -201,31 +202,31 @@ class H2DMultipoleExpansion(MultipoleExpansionBase): return list(range(-self.order, self.order+1)) def coefficients_from_source(self, avec, bvec): - from sumpy.symbolic import sympy_real_norm_2 - bessel_j = sp.Function("bessel_j") - avec_len = sympy_real_norm_2(avec) + from sumpy.symbolic import sym_real_norm_2 + bessel_j = sym.Function("bessel_j") + avec_len = sym_real_norm_2(avec) - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + k = sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) # The coordinates are negated since avec points from source to center. - source_angle_rel_center = sp.atan2(-avec[1], -avec[0]) + source_angle_rel_center = sym.atan2(-avec[1], -avec[0]) return [self.kernel.postprocess_at_source( bessel_j(l, k * avec_len) * - sp.exp(sp.I * l * -source_angle_rel_center), avec) + sym.exp(sym.I * l * -source_angle_rel_center), avec) for l in self.get_coefficient_identifiers()] def evaluate(self, coeffs, bvec): - from sumpy.symbolic import sympy_real_norm_2 - hankel_1 = sp.Function("hankel_1") - bvec_len = sympy_real_norm_2(bvec) - target_angle_rel_center = sp.atan2(bvec[1], bvec[0]) + from sumpy.symbolic import sym_real_norm_2 + hankel_1 = sym.Function("hankel_1") + bvec_len = sym_real_norm_2(bvec) + target_angle_rel_center = sym.atan2(bvec[1], bvec[0]) - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + k = sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) return sum(coeffs[self.get_storage_index(l)] * self.kernel.postprocess_at_target( hankel_1(l, k * bvec_len) - * sp.exp(sp.I * l * target_angle_rel_center), bvec) + * sym.exp(sym.I * l * target_angle_rel_center), bvec) for l in self.get_coefficient_identifiers()) def translate_from(self, src_expansion, src_coeff_exprs, dvec): @@ -233,19 +234,19 @@ class H2DMultipoleExpansion(MultipoleExpansionBase): raise RuntimeError("do not know how to translate %s to " "multipole 2D Helmholtz Bessel expansion" % type(src_expansion).__name__) - from sumpy.symbolic import sympy_real_norm_2 - dvec_len = sympy_real_norm_2(dvec) - bessel_j = sp.Function("bessel_j") - new_center_angle_rel_old_center = sp.atan2(dvec[1], dvec[0]) + from sumpy.symbolic import sym_real_norm_2 + dvec_len = sym_real_norm_2(dvec) + bessel_j = sym.Function("bessel_j") + new_center_angle_rel_old_center = sym.atan2(dvec[1], dvec[0]) - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + k = sym.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) translated_coeffs = [] for l in self.get_coefficient_identifiers(): translated_coeffs.append( sum(src_coeff_exprs[src_expansion.get_storage_index(m)] * bessel_j(m - l, k * dvec_len) - * sp.exp(sp.I * (m - l) * new_center_angle_rel_old_center) + * sym.exp(sym.I * (m - l) * new_center_angle_rel_old_center) for m in src_expansion.get_coefficient_identifiers())) return translated_coeffs diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 4a3987b5e03893f32510ae2404a7cee71b5e7639..0121cf6bcd198e0c04821d06b7110020e8428cc7 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -27,6 +27,7 @@ from six.moves import range, zip import loopy as lp import numpy as np from pymbolic.mapper import IdentityMapper, CSECachingMapperMixin +import sumpy.symbolic as sym from sumpy.symbolic import pymbolic_real_norm_2 from pymbolic.primitives import make_sym_vector from pymbolic import var @@ -263,10 +264,11 @@ class ExpressionKernel(Kernel): if self.dim != len(dist_vec): raise ValueError("dist_vec length does not match expected dimension") - expr = expr.subs([ - ("d%d" % i, dist_vec_i) + import sumpy.symbolic as sym + expr = expr.subs(dict( + (sym.Symbol("d%d" % i), dist_vec_i) for i, dist_vec_i in enumerate(dist_vec) - ]) + )) return expr @@ -506,22 +508,26 @@ 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 = ( - -var("log")(r)*(1 if icomp == jcomp else 0) - + - d[icomp]*d[jcomp]/r**2 + sum(n[axis]*d[axis] for axis in range(dim)) + * + d[icomp]*d[jcomp]/r**4 ) - scaling = 1/(4*var("pi")) + 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 = ( - (1/r)*(1 if icomp == jcomp else 0) - + - d[icomp]*d[jcomp]/r**3 + sum(n[axis]*d[axis] for axis in range(dim)) + * + d[icomp]*d[jcomp]/r**5 ) - scaling = -1/(8*var("pi")) + scaling = -3/(4*var("pi")) + elif dim is None: expr = None scaling = None @@ -558,28 +564,15 @@ 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"))] - - def postprocess_at_source(self, expr, avec): - dimensions = len(avec) - assert dimensions == self.dim - - from sumpy.symbolic import make_sympy_vector - n = make_sympy_vector(self.stresslet_vector_name, dimensions) + dim_tags="sep,C")) + ] - # 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): + def get_code_tranformer(self): from sumpy.codegen import VectorComponentRewriter vcr = VectorComponentRewriter([self.stresslet_vector_name]) from pymbolic.primitives import Variable @@ -746,12 +739,12 @@ class DirectionalTargetDerivative(DirectionalDerivative): dim = len(bvec) assert dim == self.dim - from sumpy.symbolic import make_sympy_vector - dir_vec = make_sympy_vector(self.dir_vec_name, dim) + from sumpy.symbolic import make_sym_vector + dir_vec = make_sym_vector(self.dir_vec_name, dim) # bvec = tgt-center - return sum(dir_vec[axis]*expr.diff(bvec[axis]) - for axis in range(dim)) + return sym.Add(*tuple(dir_vec[axis]*expr.diff(bvec[axis]) + for axis in range(dim))) mapper_method = "map_directional_target_derivative" @@ -777,12 +770,12 @@ class DirectionalSourceDerivative(DirectionalDerivative): dimensions = len(avec) assert dimensions == self.dim - from sumpy.symbolic import make_sympy_vector - dir_vec = make_sympy_vector(self.dir_vec_name, dimensions) + from sumpy.symbolic import make_sym_vector + dir_vec = make_sym_vector(self.dir_vec_name, dimensions) # avec = center-src -> minus sign from chain rule - return sum(-dir_vec[axis]*expr.diff(avec[axis]) - for axis in range(dimensions)) + return sym.Add(*tuple(-dir_vec[axis]*expr.diff(avec[axis]) + for axis in range(dimensions))) mapper_method = "map_directional_source_derivative" diff --git a/sumpy/p2e.py b/sumpy/p2e.py index 1b037dd4ad2ced517ddabff925ba97e6c5527348..4f687ac51e86bd44aecaeeccc5516e7a203e7c8e 100644 --- a/sumpy/p2e.py +++ b/sumpy/p2e.py @@ -70,8 +70,8 @@ class P2EBase(KernelCacheWrapper): self.dim = expansion.dim def get_loopy_instructions(self): - from sumpy.symbolic import make_sympy_vector - avec = make_sympy_vector("a", self.dim) + from sumpy.symbolic import make_sym_vector + avec = make_sym_vector("a", self.dim) from sumpy.assignment_collection import SymbolicAssignmentCollection sac = SymbolicAssignmentCollection() @@ -161,6 +161,7 @@ class P2EFromSingleBox(P2EBase): loopy_knl = self.expansion.prepare_loopy_kernel(loopy_knl) loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr") + print(loopy_knl) return loopy_knl def get_optimized_kernel(self): diff --git a/sumpy/p2p.py b/sumpy/p2p.py index 3071a3b88c9b1fc45c7d4ce3470535f084aba0a9..a8b92b91176a1ac0e0b5c9e42fc0cf71da787a6c 100644 --- a/sumpy/p2p.py +++ b/sumpy/p2p.py @@ -69,8 +69,8 @@ class P2PBase(KernelComputation, KernelCacheWrapper): self.dim = single_valued(knl.dim for knl in self.kernels) def get_loopy_insns_and_result_names(self): - from sumpy.symbolic import make_sympy_vector - dvec = make_sympy_vector("d", self.dim) + from sumpy.symbolic import make_sym_vector + dvec = make_sym_vector("d", self.dim) from sumpy.assignment_collection import SymbolicAssignmentCollection sac = SymbolicAssignmentCollection() diff --git a/sumpy/qbx.py b/sumpy/qbx.py index 0b3dc03537e788a75412b712e1c3ad18b14c11db..5012e1a668643f1bc197140eb8848d6734b09a98 100644 --- a/sumpy/qbx.py +++ b/sumpy/qbx.py @@ -27,7 +27,7 @@ import six from six.moves import range, zip import numpy as np import loopy as lp -import sympy as sp +import sumpy.symbolic as sym from pytools import memoize_method from pymbolic import parse, var @@ -64,7 +64,7 @@ def expand(expansion_nr, sac, expansion, avec, bvec): coefficients = expansion.coefficients_from_source(avec, bvec) assigned_coeffs = [ - sp.Symbol( + sym.Symbol( sac.assign_unique("expn%dcoeff%s" % ( expansion_nr, stringify_expn_index(i)), coefficients[expansion.get_storage_index(i)])) @@ -117,10 +117,10 @@ class LayerPotentialBase(KernelComputation, KernelCacheWrapper): @memoize_method def get_kernel(self): - from sumpy.symbolic import make_sympy_vector + from sumpy.symbolic import make_sym_vector - avec = make_sympy_vector("a", self.dim) - bvec = make_sympy_vector("b", self.dim) + avec = make_sym_vector("a", self.dim) + bvec = make_sym_vector("b", self.dim) from sumpy.assignment_collection import SymbolicAssignmentCollection sac = SymbolicAssignmentCollection() diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index 86c2666db8611d98c2cd6882db8c77a2abb04854..565fa8c5ad79064d69eca5e06f3ff9112c0111d2 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -27,27 +27,92 @@ import six from six.moves import range from six.moves import zip -import sympy as sp import numpy as np from pymbolic.mapper import IdentityMapper as IdentityMapperBase -from pymbolic.interop.sympy import PymbolicToSympyMapper import pymbolic.primitives as prim import logging logger = logging.getLogger(__name__) +# {{{ symbolic backend + +def _find_symbolic_backend(): + global USE_SYMENGINE + + try: + import symengine # noqa + symengine_found = True + except ImportError: + symengine_found = False + + ALLOWED_BACKENDS = ("sympy", "symengine") # noqa + + import os + backend = os.environ.get("SUMPY_FORCE_SYMBOLIC_BACKEND") + if backend is not None: + if backend not in ALLOWED_BACKENDS: + raise RuntimeError( + "SUMPY_FORCE_SYMBOLIC_BACKEND value is unrecognized: '%s' " + "(allowed values are %s)" % ( + backend, + ", ".join("'%s'" % val for val in ALLOWED_BACKENDS))) + + if backend == "symengine" and not symengine_found: + from warnings import warn + warn("SUMPY_FORCE_SYMBOLIC_BACKEND=symengine was specified, but " + "could not find symengine. Using sympy.", RuntimeWarning) + + USE_SYMENGINE = backend == "symengine" and symengine_found + else: + USE_SYMENGINE = symengine_found + + +_find_symbolic_backend() + +# }}} + +# Symbolic API common to SymEngine and sympy. +# Before adding a function here, make sure it's present in both modules. +SYMBOLIC_API = """ +Add Basic Mul Pow exp sqrt symbols sympify cos sin atan2 Function Symbol +Integer Matrix Subs I pi functions""".split() + +if USE_SYMENGINE: + from symengine import sympy_compat as sym + from pymbolic.interop.symengine import ( + PymbolicToSymEngineMapper as PymbolicToSympyMapper, + SymEngineToPymbolicMapper as SympyToPymbolicMapper) +else: + import sympy as sym + from pymbolic.interop.sympy import ( + PymbolicToSympyMapper, SympyToPymbolicMapper) + +for _apifunc in SYMBOLIC_API: + globals()[_apifunc] = getattr(sym, _apifunc) + + +def _coeff_isneg(a): + if a.is_Mul: + a = a.args[0] + try: + return a.is_Number and a < 0 + except: + return False + + # {{{ trivial assignment elimination def make_one_step_subst(assignments): - unwanted_vars = set(sp.Symbol(name) for name, value in assignments) + unwanted_vars = set(sym.Symbol(name) for name, value in assignments) result = {} - for name, value in assignments: - while value.atoms() & unwanted_vars: + assignments = dict((sym.Symbol(name), value) for name, value in assignments) + for name, value in assignments.items(): + while value.atoms(sym.Symbol) & unwanted_vars: value = value.subs(assignments) - result[sp.Symbol(name)] = value + result[name] = value return result @@ -55,12 +120,12 @@ def make_one_step_subst(assignments): def is_assignment_nontrivial(name, value): if value.is_Number: return False - elif isinstance(value, sp.Symbol): + elif isinstance(value, sym.Symbol): return False - elif (isinstance(value, sp.Mul) + elif (isinstance(value, sym.Mul) and len(value.args) == 2 and sum(1 for arg in value.args if arg.is_Number) == 1 - and sum(1 for arg in value.args if isinstance(arg, sp.Symbol)) == 1): + and sum(1 for arg in value.args if isinstance(arg, sym.Symbol)) == 1): # const*var: not good enough return False @@ -81,8 +146,10 @@ def kill_trivial_assignments(assignments, retain_names=set()): # un-substitute rejected assignments unsubst_rej = make_one_step_subst(rejected_assignments) - result = [(name, expr.xreplace(unsubst_rej)) - for name, expr in approved_assignments] + result = [] + for name, expr in approved_assignments: + r = expr.xreplace(unsubst_rej) + result.append((name, r)) logger.info("kill trivial assignments (plain): done") @@ -111,7 +178,6 @@ def _get_assignments_in_maxima(assignments, prefix=""): from pymbolic.maxima import MaximaStringifyMapper mstr = MaximaStringifyMapper() - from pymbolic.interop.sympy import SympyToPymbolicMapper s2p = SympyToPymbolicMapper() dkill = _DerivativeKiller() @@ -119,12 +185,12 @@ def _get_assignments_in_maxima(assignments, prefix=""): def write_assignment(name): symbols = [atm for atm in assignments[name].atoms() - if isinstance(atm, sp.Symbol) + if isinstance(atm, sym.Symbol) and atm.name in my_variable_names] - for sym in symbols: - if sym.name not in written_assignments: - write_assignment(sym.name) + for symb in symbols: + if symb.name not in written_assignments: + write_assignment(symb.name) result.append("%s%s : %s;" % ( prefix, name, mstr(dkill(s2p( @@ -143,7 +209,7 @@ def checked_cse(exprs, symbols=None): if symbols is not None: kwargs["symbols"] = symbols - new_assignments, new_exprs = sp.cse(exprs, **kwargs) + new_assignments, new_exprs = sym.cse(exprs, **kwargs) max_old = _get_assignments_in_maxima(dict( ("old_expr%d" % i, expr) @@ -171,8 +237,8 @@ def checked_cse(exprs, symbols=None): # }}} -def sympy_real_norm_2(x): - return sp.sqrt((x.T*x)[0, 0]) +def sym_real_norm_2(x): + return sym.sqrt((x.T*x)[0, 0]) def pymbolic_real_norm_2(x): @@ -180,9 +246,9 @@ def pymbolic_real_norm_2(x): return var("sqrt")(np.dot(x, x)) -def make_sympy_vector(name, components): - return sp.Matrix( - [sp.Symbol("%s%d" % (name, i)) for i in range(components)]) +def make_sym_vector(name, components): + return sym.Matrix( + [sym.Symbol("%s%d" % (name, i)) for i in range(components)]) def vector_subs(expr, old, new): @@ -194,8 +260,8 @@ def vector_subs(expr, old, new): def find_power_of(base, prod): - remdr = sp.Wild("remdr") - power = sp.Wild("power") + remdr = sym.Wild("remdr") + power = sym.Wild("power") result = prod.match(remdr*base**power) if result is None: return 0 @@ -205,17 +271,16 @@ def find_power_of(base, prod): class PymbolicToSympyMapperWithSymbols(PymbolicToSympyMapper): def map_variable(self, expr): if expr.name == "I": - return sp.I + return sym.I elif expr.name == "pi": - return sp.pi + return sym.pi else: return PymbolicToSympyMapper.map_variable(self, expr) def map_subscript(self, expr): if isinstance(expr.aggregate, prim.Variable) and isinstance(expr.index, int): - return sp.Symbol("%s%d" % (expr.aggregate.name, expr.index)) + return sym.Symbol("%s%d" % (expr.aggregate.name, expr.index)) else: - raise RuntimeError("do not know how to translate '%s' to sympy" - % expr) + self.raise_conversion_error(expr) # vim: fdm=marker diff --git a/sumpy/tools.py b/sumpy/tools.py index 60687510c66cd76aaca87af035464c858133d766..7d0738a750e916d0ec9678cc1fd21fa87a62c330 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -107,7 +107,7 @@ class LinearRecurrenceBasedMiDerivativeTaker(MiDerivativeTaker): expr = self.cache_by_mi[mi] except KeyError: from six import iteritems - from sympy import Add + from sumpy.symbolic import Add closest_mi = self.get_closest_cached_mi(mi) expr = self.cache_by_mi[closest_mi] @@ -264,7 +264,7 @@ class KernelComputation(object): self.name = name or self.default_name def get_kernel_scaling_assignments(self): - from pymbolic.interop.sympy import SympyToPymbolicMapper + from sumpy.symbolic import SympyToPymbolicMapper sympy_conv = SympyToPymbolicMapper() import loopy as lp diff --git a/test/test_codegen.py b/test/test_codegen.py new file mode 100644 index 0000000000000000000000000000000000000000..fcb1dd58fe614a46b604d19f136b434be0cb9e53 --- /dev/null +++ b/test/test_codegen.py @@ -0,0 +1,66 @@ +from __future__ import division, absolute_import, print_function + +__copyright__ = "Copyright (C) 2017 Matt Wala" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +import sys + +import logging +logger = logging.getLogger(__name__) + + +def test_kill_trivial_assignments(): + from sumpy.symbolic import kill_trivial_assignments, symbols, sympify + x, y, nt = symbols("x, y, nt") + t0, t1, t2 = symbols("t0:3") + u0, u1, u2 = symbols("u0:3") + + assignments = ( + ("t0", sympify(6)), + ("t1", -t0), + ("t2", 6*x), + ("nt", x**y), + # users of trivial assignments + ("u0", t0 + 1), + ("u1", t1 + 1), + ("u2", t2 + 1) + ) + + result = kill_trivial_assignments( + assignments, + retain_names=("u0", "u1", "u2")) + + assert result == [('nt', x**y), ('u0', 7), ('u1', -5), ('u2', 1 + 6*x)] + + +# You can test individual routines by typing +# $ python test_fmm.py 'test_sumpy_fmm(cl.create_some_context)' + +if __name__ == "__main__": + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from py.test.cmdline import main + main([__file__]) + +# vim: fdm=marker diff --git a/test/test_cse.py b/test/test_cse.py index d2904176e486ee16bb92bae9c8702f76969a4bae..a13ed6ab0e9d5d54563c46e299f286ab2f294f34 100644 --- a/test/test_cse.py +++ b/test/test_cse.py @@ -69,12 +69,13 @@ DAMAGE. import pytest import sys -from sympy import (Add, Pow, exp, sqrt, symbols, sympify, S, cos, sin, Eq, - Function, Tuple, CRootOf, IndexedBase, Idx, Piecewise) -from sympy.simplify.cse_opts import sub_pre, sub_post -from sympy.functions.special.hyper import meijerg -from sympy.simplify import cse_opts +from sumpy.symbolic import ( + Add, Pow, exp, sqrt, symbols, sympify, cos, sin, Function, USE_SYMENGINE) +if not USE_SYMENGINE: + from sympy.simplify.cse_opts import sub_pre, sub_post + from sympy.functions.special.hyper import meijerg + from sympy.simplify import cse_opts from sumpy.cse import ( cse, preprocess_for_cse, postprocess_for_cse) @@ -83,6 +84,9 @@ from sumpy.cse import ( w, x, y, z = symbols('w,x,y,z') x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12 = symbols('x:13') +sympyonly = ( + pytest.mark.skipif(USE_SYMENGINE, reason="uses a sympy-only feature")) + # Dummy "optimization" functions for testing. @@ -122,6 +126,7 @@ def test_cse_single(): assert reduced == [sqrt(x0) + x0**2] +@sympyonly def test_cse_not_possible(): # No substitution possible. e = Add(x, y) @@ -142,6 +147,7 @@ def test_nested_substitution(): assert reduced == [sqrt(x0) + x0**2] +@sympyonly def test_subtraction_opt(): # Make sure subtraction is optimized. e = (x - y)*(z - y) + exp((x - y)*(z - y)) @@ -195,8 +201,9 @@ def test_issue_4203(): assert cse(sin(x**x)/x**x) == ([(x0, x**x)], [sin(x0)/x0]) +@sympyonly def test_dont_cse_tuples(): - from sympy import Subs + from sumpy.symbolic import Subs f = Function("f") g = Function("g") @@ -238,9 +245,11 @@ def test_pow_invpow(): ([(x0, x**(2*y))], [x0 + 1/x0]) +@sympyonly def test_issue_4499(): # previously, this gave 16 constants from sympy.abc import a, b + from sympy import Tuple, S B = Function('B') # noqa G = Function('G') # noqa t = Tuple(* @@ -256,7 +265,9 @@ def test_issue_4499(): assert len(c[0]) == 11 +@sympyonly def test_issue_6169(): + from sympy import CRootOf r = CRootOf(x**6 - 4*x**5 - 2, 1) assert cse(r) == ([], [r]) # and a check that the right thing is done with the new @@ -264,7 +275,9 @@ def test_issue_6169(): assert sub_post(sub_pre((-x - y)*z - x - y)) == -z*(x + y) - x - y +@sympyonly def test_cse_Indexed(): # noqa + from sympy import IndexedBase, Idx len_y = 5 y = IndexedBase('y', shape=(len_y,)) x = IndexedBase('x', shape=(len_y,)) @@ -277,7 +290,9 @@ def test_cse_Indexed(): # noqa assert len(replacements) > 0 +@sympyonly def test_Piecewise(): # noqa + from sympy import Piecewise, Eq f = Piecewise((-z + x*y, Eq(y, 0)), (-z - x*y, True)) ans = cse(f) actual_ans = ([(x0, -z), (x1, x*y)], @@ -290,7 +305,7 @@ def test_name_conflict(): z2 = x2 + x3 l = [cos(z1) + z1, cos(z2) + z2, x0 + x2] substs, reduced = cse(l) - assert [e.subs(reversed(substs)) for e in reduced] == l + assert [e.subs(dict(substs)) for e in reduced] == l def test_name_conflict_cust_symbols(): @@ -298,7 +313,7 @@ def test_name_conflict_cust_symbols(): z2 = x2 + x3 l = [cos(z1) + z1, cos(z2) + z2, x0 + x2] substs, reduced = cse(l, symbols("x:10")) - assert [e.subs(reversed(substs)) for e in reduced] == l + assert [e.subs(dict(substs)) for e in reduced] == l def test_symbols_exhausted_error(): @@ -308,6 +323,7 @@ def test_symbols_exhausted_error(): print(cse(l, symbols=sym)) +@sympyonly def test_issue_7840(): # daveknippers' example C393 = sympify( # noqa