Coverage for /usr/lib/python3/dist-packages/sympy/tensor/array/array_comprehension.py: 31%

1import functools, itertools 

2from sympy.core.sympify import _sympify, sympify 

3from sympy.core.expr import Expr 

4from sympy.core import Basic, Tuple 

5from sympy.tensor.array import ImmutableDenseNDimArray 

6from sympy.core.symbol import Symbol 

7from sympy.core.numbers import Integer 

8 

9 

10class ArrayComprehension(Basic): 

11 """ 

12 Generate a list comprehension. 

13 

14 Explanation 

15 =========== 

16 

17 If there is a symbolic dimension, for example, say [i for i in range(1, N)] where 

18 N is a Symbol, then the expression will not be expanded to an array. Otherwise, 

19 calling the doit() function will launch the expansion. 

20 

21 Examples 

22 ======== 

23 

24 >>> from sympy.tensor.array import ArrayComprehension 

25 >>> from sympy import symbols 

26 >>> i, j, k = symbols('i j k') 

27 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

28 >>> a 

29 ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

30 >>> a.doit() 

31 [[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]] 

32 >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k)) 

33 >>> b.doit() 

34 ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k)) 

35 """ 

36 def __new__(cls, function, *symbols, **assumptions): 

37 if any(len(l) != 3 or None for l in symbols): 

38 raise ValueError('ArrayComprehension requires values lower and upper bound' 

39 ' for the expression') 

40 arglist = [sympify(function)] 

41 arglist.extend(cls._check_limits_validity(function, symbols)) 

42 obj = Basic.__new__(cls, *arglist, **assumptions) 

43 obj._limits = obj._args[1:] 

44 obj._shape = cls._calculate_shape_from_limits(obj._limits) 

45 obj._rank = len(obj._shape) 

46 obj._loop_size = cls._calculate_loop_size(obj._shape) 

47 return obj 

48 

49 @property 

50 def function(self): 

51 """The function applied across limits. 

52 

53 Examples 

54 ======== 

55 

56 >>> from sympy.tensor.array import ArrayComprehension 

57 >>> from sympy import symbols 

58 >>> i, j = symbols('i j') 

59 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

60 >>> a.function 

61 10*i + j 

62 """ 

63 return self._args[0] 

64 

65 @property 

66 def limits(self): 

67 """ 

68 The list of limits that will be applied while expanding the array. 

69 

70 Examples 

71 ======== 

72 

73 >>> from sympy.tensor.array import ArrayComprehension 

74 >>> from sympy import symbols 

75 >>> i, j = symbols('i j') 

76 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

77 >>> a.limits 

78 ((i, 1, 4), (j, 1, 3)) 

79 """ 

80 return self._limits 

81 

82 @property 

83 def free_symbols(self): 

84 """ 

85 The set of the free_symbols in the array. 

86 Variables appeared in the bounds are supposed to be excluded 

87 from the free symbol set. 

88 

89 Examples 

90 ======== 

91 

92 >>> from sympy.tensor.array import ArrayComprehension 

93 >>> from sympy import symbols 

94 >>> i, j, k = symbols('i j k') 

95 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

96 >>> a.free_symbols 

97 set() 

98 >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3)) 

99 >>> b.free_symbols 

100 {k} 

101 """ 

102 expr_free_sym = self.function.free_symbols 

103 for var, inf, sup in self._limits: 

104 expr_free_sym.discard(var) 

105 curr_free_syms = inf.free_symbols.union(sup.free_symbols) 

106 expr_free_sym = expr_free_sym.union(curr_free_syms) 

107 return expr_free_sym 

108 

109 @property 

110 def variables(self): 

111 """The tuples of the variables in the limits. 

112 

113 Examples 

114 ======== 

115 

116 >>> from sympy.tensor.array import ArrayComprehension 

117 >>> from sympy import symbols 

118 >>> i, j, k = symbols('i j k') 

119 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

120 >>> a.variables 

121 [i, j] 

122 """ 

123 return [l[0] for l in self._limits] 

124 

125 @property 

126 def bound_symbols(self): 

127 """The list of dummy variables. 

128 

129 Note 

130 ==== 

131 

132 Note that all variables are dummy variables since a limit without 

133 lower bound or upper bound is not accepted. 

134 """ 

135 return [l[0] for l in self._limits if len(l) != 1] 

136 

137 @property 

138 def shape(self): 

139 """ 

140 The shape of the expanded array, which may have symbols. 

141 

142 Note 

143 ==== 

144 

145 Both the lower and the upper bounds are included while 

146 calculating the shape. 

147 

148 Examples 

149 ======== 

150 

151 >>> from sympy.tensor.array import ArrayComprehension 

152 >>> from sympy import symbols 

153 >>> i, j, k = symbols('i j k') 

154 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

155 >>> a.shape 

156 (4, 3) 

157 >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3)) 

158 >>> b.shape 

159 (4, k + 3) 

160 """ 

161 return self._shape 

162 

163 @property 

164 def is_shape_numeric(self): 

165 """ 

166 Test if the array is shape-numeric which means there is no symbolic 

167 dimension. 

168 

169 Examples 

170 ======== 

171 

172 >>> from sympy.tensor.array import ArrayComprehension 

173 >>> from sympy import symbols 

174 >>> i, j, k = symbols('i j k') 

175 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

176 >>> a.is_shape_numeric 

177 True 

178 >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3)) 

179 >>> b.is_shape_numeric 

180 False 

181 """ 

182 for _, inf, sup in self._limits: 

183 if Basic(inf, sup).atoms(Symbol): 

184 return False 

185 return True 

186 

187 def rank(self): 

188 """The rank of the expanded array. 

189 

190 Examples 

191 ======== 

192 

193 >>> from sympy.tensor.array import ArrayComprehension 

194 >>> from sympy import symbols 

195 >>> i, j, k = symbols('i j k') 

196 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

197 >>> a.rank() 

198 2 

199 """ 

200 return self._rank 

201 

202 def __len__(self): 

203 """ 

204 The length of the expanded array which means the number 

205 of elements in the array. 

206 

207 Raises 

208 ====== 

209 

210 ValueError : When the length of the array is symbolic 

211 

212 Examples 

213 ======== 

214 

215 >>> from sympy.tensor.array import ArrayComprehension 

216 >>> from sympy import symbols 

217 >>> i, j = symbols('i j') 

218 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

219 >>> len(a) 

220 12 

221 """ 

222 if self._loop_size.free_symbols: 

223 raise ValueError('Symbolic length is not supported') 

224 return self._loop_size 

225 

226 @classmethod 

227 def _check_limits_validity(cls, function, limits): 

228 #limits = sympify(limits) 

229 new_limits = [] 

230 for var, inf, sup in limits: 

231 var = _sympify(var) 

232 inf = _sympify(inf) 

233 #since this is stored as an argument, it should be 

234 #a Tuple 

235 if isinstance(sup, list): 

236 sup = Tuple(*sup) 

237 else: 

238 sup = _sympify(sup) 

239 new_limits.append(Tuple(var, inf, sup)) 

240 if any((not isinstance(i, Expr)) or i.atoms(Symbol, Integer) != i.atoms() 

241 for i in [inf, sup]): 

242 raise TypeError('Bounds should be an Expression(combination of Integer and Symbol)') 

243 if (inf > sup) == True: 

244 raise ValueError('Lower bound should be inferior to upper bound') 

245 if var in inf.free_symbols or var in sup.free_symbols: 

246 raise ValueError('Variable should not be part of its bounds') 

247 return new_limits 

248 

249 @classmethod 

250 def _calculate_shape_from_limits(cls, limits): 

251 return tuple([sup - inf + 1 for _, inf, sup in limits]) 

252 

253 @classmethod 

254 def _calculate_loop_size(cls, shape): 

255 if not shape: 

256 return 0 

257 loop_size = 1 

258 for l in shape: 

259 loop_size = loop_size * l 

260 

261 return loop_size 

262 

263 def doit(self, **hints): 

264 if not self.is_shape_numeric: 

265 return self 

266 

267 return self._expand_array() 

268 

269 def _expand_array(self): 

270 res = [] 

271 for values in itertools.product(*[range(inf, sup+1) 

272 for var, inf, sup 

273 in self._limits]): 

274 res.append(self._get_element(values)) 

275 

276 return ImmutableDenseNDimArray(res, self.shape) 

277 

278 def _get_element(self, values): 

279 temp = self.function 

280 for var, val in zip(self.variables, values): 

281 temp = temp.subs(var, val) 

282 return temp 

283 

284 def tolist(self): 

285 """Transform the expanded array to a list. 

286 

287 Raises 

288 ====== 

289 

290 ValueError : When there is a symbolic dimension 

291 

292 Examples 

293 ======== 

294 

295 >>> from sympy.tensor.array import ArrayComprehension 

296 >>> from sympy import symbols 

297 >>> i, j = symbols('i j') 

298 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

299 >>> a.tolist() 

300 [[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]] 

301 """ 

302 if self.is_shape_numeric: 

303 return self._expand_array().tolist() 

304 

305 raise ValueError("A symbolic array cannot be expanded to a list") 

306 

307 def tomatrix(self): 

308 """Transform the expanded array to a matrix. 

309 

310 Raises 

311 ====== 

312 

313 ValueError : When there is a symbolic dimension 

314 ValueError : When the rank of the expanded array is not equal to 2 

315 

316 Examples 

317 ======== 

318 

319 >>> from sympy.tensor.array import ArrayComprehension 

320 >>> from sympy import symbols 

321 >>> i, j = symbols('i j') 

322 >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) 

323 >>> a.tomatrix() 

324 Matrix([ 

325 [11, 12, 13], 

326 [21, 22, 23], 

327 [31, 32, 33], 

328 [41, 42, 43]]) 

329 """ 

330 from sympy.matrices import Matrix 

331 

332 if not self.is_shape_numeric: 

333 raise ValueError("A symbolic array cannot be expanded to a matrix") 

334 if self._rank != 2: 

335 raise ValueError('Dimensions must be of size of 2') 

336 

337 return Matrix(self._expand_array().tomatrix()) 

338 

339 

340def isLambda(v): 

341 LAMBDA = lambda: 0 

342 return isinstance(v, type(LAMBDA)) and v.__name__ == LAMBDA.__name__ 

343 

344class ArrayComprehensionMap(ArrayComprehension): 

345 ''' 

346 A subclass of ArrayComprehension dedicated to map external function lambda. 

347 

348 Notes 

349 ===== 

350 

351 Only the lambda function is considered. 

352 At most one argument in lambda function is accepted in order to avoid ambiguity 

353 in value assignment. 

354 

355 Examples 

356 ======== 

357 

358 >>> from sympy.tensor.array import ArrayComprehensionMap 

359 >>> from sympy import symbols 

360 >>> i, j, k = symbols('i j k') 

361 >>> a = ArrayComprehensionMap(lambda: 1, (i, 1, 4)) 

362 >>> a.doit() 

363 [1, 1, 1, 1] 

364 >>> b = ArrayComprehensionMap(lambda a: a+1, (j, 1, 4)) 

365 >>> b.doit() 

366 [2, 3, 4, 5] 

367 

368 ''' 

369 def __new__(cls, function, *symbols, **assumptions): 

370 if any(len(l) != 3 or None for l in symbols): 

371 raise ValueError('ArrayComprehension requires values lower and upper bound' 

372 ' for the expression') 

373 

374 if not isLambda(function): 

375 raise ValueError('Data type not supported') 

376 

377 arglist = cls._check_limits_validity(function, symbols) 

378 obj = Basic.__new__(cls, *arglist, **assumptions) 

379 obj._limits = obj._args 

380 obj._shape = cls._calculate_shape_from_limits(obj._limits) 

381 obj._rank = len(obj._shape) 

382 obj._loop_size = cls._calculate_loop_size(obj._shape) 

383 obj._lambda = function 

384 return obj 

385 

386 @property 

387 def func(self): 

388 class _(ArrayComprehensionMap): 

389 def __new__(cls, *args, **kwargs): 

390 return ArrayComprehensionMap(self._lambda, *args, **kwargs) 

391 return _ 

392 

393 def _get_element(self, values): 

394 temp = self._lambda 

395 if self._lambda.__code__.co_argcount == 0: 

396 temp = temp() 

397 elif self._lambda.__code__.co_argcount == 1: 

398 temp = temp(functools.reduce(lambda a, b: a*b, values)) 

399 return temp