Coverage for /usr/lib/python3/dist-packages/sympy/simplify/cse_opts.py: 79%

1""" Optimizations of the expression tree representation for better CSE 

2opportunities. 

3""" 

4from sympy.core import Add, Basic, Mul 

5from sympy.core.singleton import S 

6from sympy.core.sorting import default_sort_key 

7from sympy.core.traversal import preorder_traversal 

8 

9 

10def sub_pre(e): 

11 """ Replace y - x with -(x - y) if -1 can be extracted from y - x. 

12 """ 

13 # replacing Add, A, from which -1 can be extracted with -1*-A 

14 adds = [a for a in e.atoms(Add) if a.could_extract_minus_sign()] 

15 reps = {} 

16 ignore = set() 

17 for a in adds: 

18 na = -a 

19 if na.is_Mul: # e.g. MatExpr 

20 ignore.add(a) 

21 continue 

22 reps[a] = Mul._from_args([S.NegativeOne, na]) 

23 

24 e = e.xreplace(reps) 

25 

26 # repeat again for persisting Adds but mark these with a leading 1, -1 

27 # e.g. y - x -> 1*-1*(x - y) 

28 if isinstance(e, Basic): 

29 negs = {} 

30 for a in sorted(e.atoms(Add), key=default_sort_key): 

31 if a in ignore: 

32 continue 

33 if a in reps: 

34 negs[a] = reps[a] 

35 elif a.could_extract_minus_sign(): 

36 negs[a] = Mul._from_args([S.One, S.NegativeOne, -a]) 

37 e = e.xreplace(negs) 

38 return e 

39 

40 

41def sub_post(e): 

42 """ Replace 1*-1*x with -x. 

43 """ 

44 replacements = [] 

45 for node in preorder_traversal(e): 

46 if isinstance(node, Mul) and \ 

47 node.args[0] is S.One and node.args[1] is S.NegativeOne: 

48 replacements.append((node, -Mul._from_args(node.args[2:]))) 

49 for node, replacement in replacements: 

50 e = e.xreplace({node: replacement}) 

51 

52 return e