1""" 

2Default Sympy optimizations applied in pystencils kernels using :func:`sympy.codegen.rewriting.optimize`. 

3 

4See :func:`sympy.codegen.rewriting.optimize`. 

5""" 

6 

7 

8import itertools 

9 

10from pystencils import Assignment 

11from pystencils.astnodes import SympyAssignment 

12 

13try: 

14 from sympy.codegen.rewriting import optims_c99, optimize 

15 from sympy.codegen.rewriting import ReplaceOptim 

16 HAS_REWRITING = True 

17 

18 # Evaluates all constant terms 

19 evaluate_constant_terms = ReplaceOptim( 

20 lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer, 

21 lambda p: p.evalf() 

22 ) 

23 

24 optims_pystencils_cpu = [evaluate_constant_terms] + list(optims_c99) 

25 optims_pystencils_gpu = [evaluate_constant_terms] + list(optims_c99) 

26except ImportError: 

27 from warnings import warn 

28 warn("Could not import ReplaceOptim, optims_c99, optimize from sympy.codegen.rewriting." 

29 "Please update your sympy installation!") 

30 optims_c99 = [] 

31 optims_pystencils_cpu = [] 

32 optims_pystencils_gpu = [] 

33 HAS_REWRITING = False 

34 

35 

36def optimize_assignments(assignments, optimizations): 

37 

38 if HAS_REWRITING: 

39 assignments = [Assignment(a.lhs, optimize(a.rhs, optimizations)) 

40 if hasattr(a, 'lhs') 

41 else a for a in assignments] 

42 assignments_nodes = [a.atoms(SympyAssignment) for a in assignments] 

43 for a in itertools.chain.from_iterable(assignments_nodes): 

44 a.optimize(optimizations) 

45 

46 return assignments 

47 

48 

49def optimize_ast(ast, optimizations): 

50 

51 if HAS_REWRITING: 

52 assignments_nodes = ast.atoms(SympyAssignment) 

53 for a in assignments_nodes: 

54 a.optimize(optimizations) 

55 

56 return ast