def test_replace_powers_with_gray_functions(self): gray_square = sympy.Function('GRAY_SQUARE') gray_cube = sympy.Function('GRAY_CUBE') gray_four = sympy.Function('GRAY_FOUR') gray_sqrt = sympy.Function('GRAY_SQRT') gray_sqrt_cube = sympy.Function('GRAY_SQRT_CUBE') repl = { 2: gray_square, 3: gray_cube, 4: gray_four, 0.5: gray_sqrt, 1.5: gray_sqrt_cube } a, b, c = sympy.symbols('a b c') expand_opt4 = create_expand_pow_optimization(4) with self.assertRaises(TypeError): codegen.replace_powers_with_gray_functions(0, repl) self.assertEqual( codegen.replace_powers_with_gray_functions(a**2, repl), expand_opt4(a * a)) self.assertEqual( codegen.replace_powers_with_gray_functions((a + b)**2, repl), gray_square(a + b)) self.assertEqual( codegen.replace_powers_with_gray_functions(1 / (a + b)**2, repl), 1 / gray_square(a + b)) self.assertEqual( codegen.replace_powers_with_gray_functions(1 / (a + b)**2 + a**5, repl), 1 / gray_square(a + b) + a**5) self.assertEqual( codegen.replace_powers_with_gray_functions( 1 / (a + b)**2 + a**5 + sympy.sqrt(b), repl), 1 / gray_square(a + b) + a**5 + gray_sqrt(b)) self.assertEqual( codegen.replace_powers_with_gray_functions( 1 / (a + b)**2 + a**5 + sympy.sqrt(a + b + c), repl), 1 / gray_square(a + b) + a**5 + gray_sqrt(a + b + c)) self.assertEqual( codegen.replace_powers_with_gray_functions(1 / (a + b)**(3 / 2), repl), 1 / gray_sqrt_cube(a + b)) self.assertEqual( codegen.replace_powers_with_gray_functions(1 / (a + b)**(1.5), repl), 1 / gray_sqrt_cube(a + b))
def __init__(self, stencil, dimensions): self.stencil = stencil self.dimensions = dimensions self.time_access_pattern = re.compile('u\[(.*?)\]') #self.pow_pattern = re.compile('pow\(.*?\)') self.expand_opt = create_expand_pow_optimization(2)
def test_create_expand_pow_optimization(): cc = lambda x: ccode( optimize(x, [create_expand_pow_optimization(4)])) x = Symbol('x') assert cc(x**4) == 'x*x*x*x' assert cc(x**4 + x**2) == 'x*x + x*x*x*x' assert cc(x**5 + x**4) == 'pow(x, 5) + x*x*x*x' assert cc(sin(x)**4) == 'pow(sin(x), 4)' # gh issue 15335 assert cc(x**(-4)) == '1.0/(x*x*x*x)' assert cc(x**(-5)) == 'pow(x, -5)' assert cc(-x**4) == '-x*x*x*x' assert cc(x**4 - x**2) == '-x*x + x*x*x*x' i = Symbol('i', integer=True) assert cc(x**i - x**2) == 'pow(x, i) - x*x' # gh issue 20753 cc2 = lambda x: ccode(optimize(x, [create_expand_pow_optimization( 4, base_req=lambda b: b.is_Function)])) assert cc2(x**3 + sin(x)**3) == "pow(x, 3) + sin(x)*sin(x)*sin(x)"
def test_create_expand_pow_optimization(): my_opt = create_expand_pow_optimization(4) x = Symbol('x') assert ccode(optimize(x**4, [my_opt])) == 'x*x*x*x' x5x4 = x**5 + x**4 assert ccode(optimize(x5x4, [my_opt])) == 'pow(x, 5) + x*x*x*x' sin4x = sin(x)**4 assert ccode(optimize(sin4x, [my_opt])) == 'pow(sin(x), 4)'
def test_create_expand_pow_optimization(): my_opt = create_expand_pow_optimization(4) x = Symbol('x') assert ccode(optimize(x**4, [my_opt])) == 'x*x*x*x' x5x4 = x**5 + x**4 assert ccode(optimize(x5x4, [my_opt])) == 'pow(x, 5) + x*x*x*x' sin4x = sin(x)**4 assert ccode(optimize(sin4x, [my_opt])) == 'pow(sin(x), 4)' assert ccode(optimize((x**(-4)), [my_opt])) == 'pow(x, -4)'
def test_create_expand_pow_optimization(): cc = lambda x: ccode(optimize(x, [create_expand_pow_optimization(4)])) x = Symbol('x') assert cc(x**4) == 'x*x*x*x' assert cc(x**4 + x**2) == 'x*x + x*x*x*x' assert cc(x**5 + x**4) == 'pow(x, 5) + x*x*x*x' assert cc(sin(x)**4) == 'pow(sin(x), 4)' # gh issue 15335 assert cc(x**(-4)) == '1.0/(x*x*x*x)' assert cc(x**(-5)) == 'pow(x, -5)' assert cc(-x**4) == '-x*x*x*x' assert cc(x**4 - x**2) == '-x*x + x*x*x*x' i = Symbol('i', integer=True) assert cc(x**i - x**2) == 'pow(x, i) - x*x'
def test_create_expand_pow_optimization(): cc = lambda x: ccode(optimize(x, [create_expand_pow_optimization(4)])) x = Symbol("x") assert cc(x**4) == "x*x*x*x" assert cc(x**4 + x**2) == "x*x + x*x*x*x" assert cc(x**5 + x**4) == "pow(x, 5) + x*x*x*x" assert cc(sin(x)**4) == "pow(sin(x), 4)" # gh issue 15335 assert cc(x**(-4)) == "1.0/(x*x*x*x)" assert cc(x**(-5)) == "pow(x, -5)" assert cc(-(x**4)) == "-x*x*x*x" assert cc(x**4 - x**2) == "-x*x + x*x*x*x" i = Symbol("i", integer=True) assert cc(x**i - x**2) == "pow(x, i) - x*x"
def replace_powers_with_gray_functions(expr, repl, up_to_exponent=4): """repls is a dictionary of the form: {exponent_of_the_power: sympy.Function("GRAY_FUNCTION")}. GRAY_FUNCTIONS have to be defined in GRay! The definition can be added to the template of the output file. """ if (not isinstance(expr, Basic)): raise TypeError( "replace_powers_with_gray_functions can only handle sympy expressions!" ) # Instead of having pow(x,3), write x*x*x # create_expand_pow_optimization(N) performs this substitution # up to the power N expand_opt = create_expand_pow_optimization(up_to_exponent) expr = expand_opt(expr) # The previous substitution does not work when there are multiple # symbols involved, for example (a+b)**2 # Having performed cse, these expressions should not be too bad a, b = Wild('a'), Wild('b') for e in repl.keys(): expr = expr.replace((a + b)**e, lambda a, b: repl[e](a + b)) expr = expr.replace((a + b)**-e, lambda a, b: 1 / repl[e](a + b)) # For the square roots we also match single symbols # Square roots are identified by non-integer exponents sqrt_elems = [e for e in repl.keys() if type(e) is float] for e in sqrt_elems: expr = expr.replace(a**e, lambda a: repl[e](a)) expr = expr.replace(a**-e, lambda a: 1 / repl[e](a)) return expr
def csetupletocline(t, real_type): """Returns a string line of c code of t, a tuple from sympy cse. Uses real_type a string of either "double" or "float" """ expand_opt = create_expand_pow_optimization(3) return real_type + " " + sp.ccode(Assignment(t[0], expand_opt(t[1])))
#!/usr/bin/env python import sympy import sys from sympy.codegen.rewriting import create_expand_pow_optimization maxorder = int(sys.argv[1]) maxorder_damped = int(sys.argv[2]) maxpow = int(sys.argv[3]) expand_pow = create_expand_pow_optimization(maxpow) print( "// maximum order, maximum order for Thole-damped tensors, maximum expansion order of pow" ) print( f"// maxorder = {maxorder}, maxorder_damped = {maxorder_damped}, maxpow = {maxpow}" ) def T_cart(sequence): """ T = 1/R Generates sympy cartesian T-tensor of rank len(sequence) """ x, y, z = sympy.symbols("x y z") R = sympy.sqrt(x**2 + y**2 + z**2) T = 1 / R if sequence: return sympy.diff(T, *[i for i in sequence]) else: return T