1import numpy as np 

2import sympy as sp 

3 

4from pystencils.data_types import cast_func, collate_types, create_type, get_type_of_expression 

5from pystencils.sympyextensions import is_integer_sequence 

6 

7 

8class IntegerFunctionTwoArgsMixIn(sp.Function): 

9 is_integer = True 

10 

11 def __new__(cls, arg1, arg2): 

12 args = [] 

13 for a in (arg1, arg2): 

14 if isinstance(a, sp.Number) or isinstance(a, int): 

15 args.append(cast_func(a, create_type("int"))) 

16 elif isinstance(a, np.generic): 

17 args.append(cast_func(a, a.dtype)) 

18 else: 

19 args.append(a) 

20 

21 for a in args: 

22 try: 

23 type = get_type_of_expression(a) 

24 if not type.is_int(): 

25 raise ValueError("Argument to integer function is not an int but " + str(type)) 

26 except NotImplementedError: 

27 raise ValueError("Integer functions can only be constructed with typed expressions") 

28 return super().__new__(cls, *args) 

29 

30 

31# noinspection PyPep8Naming 

32class bitwise_xor(IntegerFunctionTwoArgsMixIn): 

33 pass 

34 

35 

36# noinspection PyPep8Naming 

37class bit_shift_right(IntegerFunctionTwoArgsMixIn): 

38 pass 

39 

40 

41# noinspection PyPep8Naming 

42class bit_shift_left(IntegerFunctionTwoArgsMixIn): 

43 pass 

44 

45 

46# noinspection PyPep8Naming 

47class bitwise_and(IntegerFunctionTwoArgsMixIn): 

48 pass 

49 

50 

51# noinspection PyPep8Naming 

52class bitwise_or(IntegerFunctionTwoArgsMixIn): 

53 pass 

54 

55 

56# noinspection PyPep8Naming 

57class int_div(IntegerFunctionTwoArgsMixIn): 

58 pass 

59 

60 

61# noinspection PyPep8Naming 

62class int_power_of_2(IntegerFunctionTwoArgsMixIn): 

63 pass 

64 

65 

66# noinspection PyPep8Naming 

67class modulo_floor(sp.Function): 

68 """Returns the next smaller integer divisible by given divisor. 

69 

70 Examples: 

71 >>> modulo_floor(9, 4) 

72 8 

73 >>> modulo_floor(11, 4) 

74 8 

75 >>> modulo_floor(12, 4) 

76 12 

77 >>> from pystencils import TypedSymbol 

78 >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32") 

79 >>> modulo_floor(a, b).to_c(str) 

80 '(int64_t)((a) / (b)) * (b)' 

81 """ 

82 nargs = 2 

83 is_integer = True 

84 

85 def __new__(cls, integer, divisor): 

86 if is_integer_sequence((integer, divisor)): 

87 return (int(integer) // int(divisor)) * divisor 

88 else: 

89 return super().__new__(cls, integer, divisor) 

90 

91 def to_c(self, print_func): 

92 dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1]))) 

93 assert dtype.is_int() 

94 return "({dtype})(({0}) / ({1})) * ({1})".format(print_func(self.args[0]), 

95 print_func(self.args[1]), dtype=dtype) 

96 

97 

98# noinspection PyPep8Naming 

99class modulo_ceil(sp.Function): 

100 """Returns the next bigger integer divisible by given divisor. 

101 

102 Examples: 

103 >>> modulo_ceil(9, 4) 

104 12 

105 >>> modulo_ceil(11, 4) 

106 12 

107 >>> modulo_ceil(12, 4) 

108 12 

109 >>> from pystencils import TypedSymbol 

110 >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32") 

111 >>> modulo_ceil(a, b).to_c(str) 

112 '((a) % (b) == 0 ? a : ((int64_t)((a) / (b))+1) * (b))' 

113 """ 

114 nargs = 2 

115 is_integer = True 

116 

117 def __new__(cls, integer, divisor): 

118 if is_integer_sequence((integer, divisor)): 118 ↛ 121line 118 didn't jump to line 121, because the condition on line 118 was never false

119 return integer if integer % divisor == 0 else ((integer // divisor) + 1) * divisor 

120 else: 

121 return super().__new__(cls, integer, divisor) 

122 

123 def to_c(self, print_func): 

124 dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1]))) 

125 assert dtype.is_int() 

126 code = "(({0}) % ({1}) == 0 ? {0} : (({dtype})(({0}) / ({1}))+1) * ({1}))" 

127 return code.format(print_func(self.args[0]), print_func(self.args[1]), dtype=dtype) 

128 

129 

130# noinspection PyPep8Naming 

131class div_ceil(sp.Function): 

132 """Integer division that is always rounded up 

133 

134 Examples: 

135 >>> div_ceil(9, 4) 

136 3 

137 >>> div_ceil(8, 4) 

138 2 

139 >>> from pystencils import TypedSymbol 

140 >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32") 

141 >>> div_ceil(a, b).to_c(str) 

142 '( (a) % (b) == 0 ? (int64_t)(a) / (int64_t)(b) : ( (int64_t)(a) / (int64_t)(b) ) +1 )' 

143 """ 

144 nargs = 2 

145 is_integer = True 

146 

147 def __new__(cls, integer, divisor): 

148 if is_integer_sequence((integer, divisor)): 

149 return integer // divisor if integer % divisor == 0 else (integer // divisor) + 1 

150 else: 

151 return super().__new__(cls, integer, divisor) 

152 

153 def to_c(self, print_func): 

154 dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1]))) 

155 assert dtype.is_int() 

156 code = "( ({0}) % ({1}) == 0 ? ({dtype})({0}) / ({dtype})({1}) : ( ({dtype})({0}) / ({dtype})({1}) ) +1 )" 

157 return code.format(print_func(self.args[0]), print_func(self.args[1]), dtype=dtype) 

158 

159 

160# noinspection PyPep8Naming 

161class div_floor(sp.Function): 

162 """Integer division 

163 

164 Examples: 

165 >>> div_floor(9, 4) 

166 2 

167 >>> div_floor(8, 4) 

168 2 

169 >>> from pystencils import TypedSymbol 

170 >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32") 

171 >>> div_floor(a, b).to_c(str) 

172 '((int64_t)(a) / (int64_t)(b))' 

173 """ 

174 nargs = 2 

175 is_integer = True 

176 

177 def __new__(cls, integer, divisor): 

178 if is_integer_sequence((integer, divisor)): 

179 return integer // divisor 

180 else: 

181 return super().__new__(cls, integer, divisor) 

182 

183 def to_c(self, print_func): 

184 dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1]))) 

185 assert dtype.is_int() 

186 code = "(({dtype})({0}) / ({dtype})({1}))" 

187 return code.format(print_func(self.args[0]), print_func(self.args[1]), dtype=dtype)