Coverage for /usr/lib/python3/dist-packages/sympy/strategies/rl.py: 30%

1""" Generic Rules for SymPy 

2 

3This file assumes knowledge of Basic and little else. 

4""" 

5from sympy.utilities.iterables import sift 

6from .util import new 

7 

8 

9# Functions that create rules 

10def rm_id(isid, new=new): 

11 """ Create a rule to remove identities. 

12 

13 isid - fn :: x -> Bool --- whether or not this element is an identity. 

14 

15 Examples 

16 ======== 

17 

18 >>> from sympy.strategies import rm_id 

19 >>> from sympy import Basic, S 

20 >>> remove_zeros = rm_id(lambda x: x==0) 

21 >>> remove_zeros(Basic(S(1), S(0), S(2))) 

22 Basic(1, 2) 

23 >>> remove_zeros(Basic(S(0), S(0))) # If only identites then we keep one 

24 Basic(0) 

25 

26 See Also: 

27 unpack 

28 """ 

29 def ident_remove(expr): 

30 """ Remove identities """ 

31 ids = list(map(isid, expr.args)) 

32 if sum(ids) == 0: # No identities. Common case 

33 return expr 

34 elif sum(ids) != len(ids): # there is at least one non-identity 

35 return new(expr.__class__, 

36 *[arg for arg, x in zip(expr.args, ids) if not x]) 

37 else: 

38 return new(expr.__class__, expr.args[0]) 

39 

40 return ident_remove 

41 

42 

43def glom(key, count, combine): 

44 """ Create a rule to conglomerate identical args. 

45 

46 Examples 

47 ======== 

48 

49 >>> from sympy.strategies import glom 

50 >>> from sympy import Add 

51 >>> from sympy.abc import x 

52 

53 >>> key = lambda x: x.as_coeff_Mul()[1] 

54 >>> count = lambda x: x.as_coeff_Mul()[0] 

55 >>> combine = lambda cnt, arg: cnt * arg 

56 >>> rl = glom(key, count, combine) 

57 

58 >>> rl(Add(x, -x, 3*x, 2, 3, evaluate=False)) 

59 3*x + 5 

60 

61 Wait, how are key, count and combine supposed to work? 

62 

63 >>> key(2*x) 

64 x 

65 >>> count(2*x) 

66 2 

67 >>> combine(2, x) 

68 2*x 

69 """ 

70 def conglomerate(expr): 

71 """ Conglomerate together identical args x + x -> 2x """ 

72 groups = sift(expr.args, key) 

73 counts = {k: sum(map(count, args)) for k, args in groups.items()} 

74 newargs = [combine(cnt, mat) for mat, cnt in counts.items()] 

75 if set(newargs) != set(expr.args): 

76 return new(type(expr), *newargs) 

77 else: 

78 return expr 

79 

80 return conglomerate 

81 

82 

83def sort(key, new=new): 

84 """ Create a rule to sort by a key function. 

85 

86 Examples 

87 ======== 

88 

89 >>> from sympy.strategies import sort 

90 >>> from sympy import Basic, S 

91 >>> sort_rl = sort(str) 

92 >>> sort_rl(Basic(S(3), S(1), S(2))) 

93 Basic(1, 2, 3) 

94 """ 

95 

96 def sort_rl(expr): 

97 return new(expr.__class__, *sorted(expr.args, key=key)) 

98 return sort_rl 

99 

100 

101def distribute(A, B): 

102 """ Turns an A containing Bs into a B of As 

103 

104 where A, B are container types 

105 

106 >>> from sympy.strategies import distribute 

107 >>> from sympy import Add, Mul, symbols 

108 >>> x, y = symbols('x,y') 

109 >>> dist = distribute(Mul, Add) 

110 >>> expr = Mul(2, x+y, evaluate=False) 

111 >>> expr 

112 2*(x + y) 

113 >>> dist(expr) 

114 2*x + 2*y 

115 """ 

116 

117 def distribute_rl(expr): 

118 for i, arg in enumerate(expr.args): 

119 if isinstance(arg, B): 

120 first, b, tail = expr.args[:i], expr.args[i], expr.args[i + 1:] 

121 return B(*[A(*(first + (arg,) + tail)) for arg in b.args]) 

122 return expr 

123 return distribute_rl 

124 

125 

126def subs(a, b): 

127 """ Replace expressions exactly """ 

128 def subs_rl(expr): 

129 if expr == a: 

130 return b 

131 else: 

132 return expr 

133 return subs_rl 

134 

135 

136# Functions that are rules 

137def unpack(expr): 

138 """ Rule to unpack singleton args 

139 

140 >>> from sympy.strategies import unpack 

141 >>> from sympy import Basic, S 

142 >>> unpack(Basic(S(2))) 

143 2 

144 """ 

145 if len(expr.args) == 1: 

146 return expr.args[0] 

147 else: 

148 return expr 

149 

150 

151def flatten(expr, new=new): 

152 """ Flatten T(a, b, T(c, d), T2(e)) to T(a, b, c, d, T2(e)) """ 

153 cls = expr.__class__ 

154 args = [] 

155 for arg in expr.args: 

156 if arg.__class__ == cls: 

157 args.extend(arg.args) 

158 else: 

159 args.append(arg) 

160 return new(expr.__class__, *args) 

161 

162 

163def rebuild(expr): 

164 """ Rebuild a SymPy tree. 

165 

166 Explanation 

167 =========== 

168 

169 This function recursively calls constructors in the expression tree. 

170 This forces canonicalization and removes ugliness introduced by the use of 

171 Basic.__new__ 

172 """ 

173 if expr.is_Atom: 

174 return expr 

175 else: 

176 return expr.func(*list(map(rebuild, expr.args)))