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
