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

1import itertools 

2from collections.abc import Iterable 

3 

4from sympy.core._print_helpers import Printable 

5from sympy.core.containers import Tuple 

6from sympy.core.function import diff 

7from sympy.core.singleton import S 

8from sympy.core.sympify import _sympify 

9 

10from sympy.tensor.array.ndim_array import NDimArray 

11from sympy.tensor.array.dense_ndim_array import DenseNDimArray, ImmutableDenseNDimArray 

12from sympy.tensor.array.sparse_ndim_array import SparseNDimArray 

13 

14 

15def _arrayfy(a): 

16 from sympy.matrices import MatrixBase 

17 

18 if isinstance(a, NDimArray): 

19 return a 

20 if isinstance(a, (MatrixBase, list, tuple, Tuple)): 

21 return ImmutableDenseNDimArray(a) 

22 return a 

23 

24 

25def tensorproduct(*args): 

26 """ 

27 Tensor product among scalars or array-like objects. 

28 

29 The equivalent operator for array expressions is ``ArrayTensorProduct``, 

30 which can be used to keep the expression unevaluated. 

31 

32 Examples 

33 ======== 

34 

35 >>> from sympy.tensor.array import tensorproduct, Array 

36 >>> from sympy.abc import x, y, z, t 

37 >>> A = Array([[1, 2], [3, 4]]) 

38 >>> B = Array([x, y]) 

39 >>> tensorproduct(A, B) 

40 [[[x, y], [2*x, 2*y]], [[3*x, 3*y], [4*x, 4*y]]] 

41 >>> tensorproduct(A, x) 

42 [[x, 2*x], [3*x, 4*x]] 

43 >>> tensorproduct(A, B, B) 

44 [[[[x**2, x*y], [x*y, y**2]], [[2*x**2, 2*x*y], [2*x*y, 2*y**2]]], [[[3*x**2, 3*x*y], [3*x*y, 3*y**2]], [[4*x**2, 4*x*y], [4*x*y, 4*y**2]]]] 

45 

46 Applying this function on two matrices will result in a rank 4 array. 

47 

48 >>> from sympy import Matrix, eye 

49 >>> m = Matrix([[x, y], [z, t]]) 

50 >>> p = tensorproduct(eye(3), m) 

51 >>> p 

52 [[[[x, y], [z, t]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[x, y], [z, t]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[x, y], [z, t]]]] 

53 

54 See Also 

55 ======== 

56 

57 sympy.tensor.array.expressions.array_expressions.ArrayTensorProduct 

58 

59 """ 

60 from sympy.tensor.array import SparseNDimArray, ImmutableSparseNDimArray 

61 

62 if len(args) == 0: 

63 return S.One 

64 if len(args) == 1: 

65 return _arrayfy(args[0]) 

66 from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract 

67 from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct 

68 from sympy.tensor.array.expressions.array_expressions import _ArrayExpr 

69 from sympy.matrices.expressions.matexpr import MatrixSymbol 

70 if any(isinstance(arg, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)) for arg in args): 

71 return ArrayTensorProduct(*args) 

72 if len(args) > 2: 

73 return tensorproduct(tensorproduct(args[0], args[1]), *args[2:]) 

74 

75 # length of args is 2: 

76 a, b = map(_arrayfy, args) 

77 

78 if not isinstance(a, NDimArray) or not isinstance(b, NDimArray): 

79 return a*b 

80 

81 if isinstance(a, SparseNDimArray) and isinstance(b, SparseNDimArray): 

82 lp = len(b) 

83 new_array = {k1*lp + k2: v1*v2 for k1, v1 in a._sparse_array.items() for k2, v2 in b._sparse_array.items()} 

84 return ImmutableSparseNDimArray(new_array, a.shape + b.shape) 

85 

86 product_list = [i*j for i in Flatten(a) for j in Flatten(b)] 

87 return ImmutableDenseNDimArray(product_list, a.shape + b.shape) 

88 

89 

90def _util_contraction_diagonal(array, *contraction_or_diagonal_axes): 

91 array = _arrayfy(array) 

92 

93 # Verify contraction_axes: 

94 taken_dims = set() 

95 for axes_group in contraction_or_diagonal_axes: 

96 if not isinstance(axes_group, Iterable): 

97 raise ValueError("collections of contraction/diagonal axes expected") 

98 

99 dim = array.shape[axes_group[0]] 

100 

101 for d in axes_group: 

102 if d in taken_dims: 

103 raise ValueError("dimension specified more than once") 

104 if dim != array.shape[d]: 

105 raise ValueError("cannot contract or diagonalize between axes of different dimension") 

106 taken_dims.add(d) 

107 

108 rank = array.rank() 

109 

110 remaining_shape = [dim for i, dim in enumerate(array.shape) if i not in taken_dims] 

111 cum_shape = [0]*rank 

112 _cumul = 1 

113 for i in range(rank): 

114 cum_shape[rank - i - 1] = _cumul 

115 _cumul *= int(array.shape[rank - i - 1]) 

116 

117 # DEFINITION: by absolute position it is meant the position along the one 

118 # dimensional array containing all the tensor components. 

119 

120 # Possible future work on this module: move computation of absolute 

121 # positions to a class method. 

122 

123 # Determine absolute positions of the uncontracted indices: 

124 remaining_indices = [[cum_shape[i]*j for j in range(array.shape[i])] 

125 for i in range(rank) if i not in taken_dims] 

126 

127 # Determine absolute positions of the contracted indices: 

128 summed_deltas = [] 

129 for axes_group in contraction_or_diagonal_axes: 

130 lidx = [] 

131 for js in range(array.shape[axes_group[0]]): 

132 lidx.append(sum([cum_shape[ig] * js for ig in axes_group])) 

133 summed_deltas.append(lidx) 

134 

135 return array, remaining_indices, remaining_shape, summed_deltas 

136 

137 

138def tensorcontraction(array, *contraction_axes): 

139 """ 

140 Contraction of an array-like object on the specified axes. 

141 

142 The equivalent operator for array expressions is ``ArrayContraction``, 

143 which can be used to keep the expression unevaluated. 

144 

145 Examples 

146 ======== 

147 

148 >>> from sympy import Array, tensorcontraction 

149 >>> from sympy import Matrix, eye 

150 >>> tensorcontraction(eye(3), (0, 1)) 

151 3 

152 >>> A = Array(range(18), (3, 2, 3)) 

153 >>> A 

154 [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]] 

155 >>> tensorcontraction(A, (0, 2)) 

156 [21, 30] 

157 

158 Matrix multiplication may be emulated with a proper combination of 

159 ``tensorcontraction`` and ``tensorproduct`` 

160 

161 >>> from sympy import tensorproduct 

162 >>> from sympy.abc import a,b,c,d,e,f,g,h 

163 >>> m1 = Matrix([[a, b], [c, d]]) 

164 >>> m2 = Matrix([[e, f], [g, h]]) 

165 >>> p = tensorproduct(m1, m2) 

166 >>> p 

167 [[[[a*e, a*f], [a*g, a*h]], [[b*e, b*f], [b*g, b*h]]], [[[c*e, c*f], [c*g, c*h]], [[d*e, d*f], [d*g, d*h]]]] 

168 >>> tensorcontraction(p, (1, 2)) 

169 [[a*e + b*g, a*f + b*h], [c*e + d*g, c*f + d*h]] 

170 >>> m1*m2 

171 Matrix([ 

172 [a*e + b*g, a*f + b*h], 

173 [c*e + d*g, c*f + d*h]]) 

174 

175 See Also 

176 ======== 

177 

178 sympy.tensor.array.expressions.array_expressions.ArrayContraction 

179 

180 """ 

181 from sympy.tensor.array.expressions.array_expressions import _array_contraction 

182 from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract 

183 from sympy.tensor.array.expressions.array_expressions import _ArrayExpr 

184 from sympy.matrices.expressions.matexpr import MatrixSymbol 

185 if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)): 

186 return _array_contraction(array, *contraction_axes) 

187 

188 array, remaining_indices, remaining_shape, summed_deltas = _util_contraction_diagonal(array, *contraction_axes) 

189 

190 # Compute the contracted array: 

191 # 

192 # 1. external for loops on all uncontracted indices. 

193 # Uncontracted indices are determined by the combinatorial product of 

194 # the absolute positions of the remaining indices. 

195 # 2. internal loop on all contracted indices. 

196 # It sums the values of the absolute contracted index and the absolute 

197 # uncontracted index for the external loop. 

198 contracted_array = [] 

199 for icontrib in itertools.product(*remaining_indices): 

200 index_base_position = sum(icontrib) 

201 isum = S.Zero 

202 for sum_to_index in itertools.product(*summed_deltas): 

203 idx = array._get_tuple_index(index_base_position + sum(sum_to_index)) 

204 isum += array[idx] 

205 

206 contracted_array.append(isum) 

207 

208 if len(remaining_indices) == 0: 

209 assert len(contracted_array) == 1 

210 return contracted_array[0] 

211 

212 return type(array)(contracted_array, remaining_shape) 

213 

214 

215def tensordiagonal(array, *diagonal_axes): 

216 """ 

217 Diagonalization of an array-like object on the specified axes. 

218 

219 This is equivalent to multiplying the expression by Kronecker deltas 

220 uniting the axes. 

221 

222 The diagonal indices are put at the end of the axes. 

223 

224 The equivalent operator for array expressions is ``ArrayDiagonal``, which 

225 can be used to keep the expression unevaluated. 

226 

227 Examples 

228 ======== 

229 

230 ``tensordiagonal`` acting on a 2-dimensional array by axes 0 and 1 is 

231 equivalent to the diagonal of the matrix: 

232 

233 >>> from sympy import Array, tensordiagonal 

234 >>> from sympy import Matrix, eye 

235 >>> tensordiagonal(eye(3), (0, 1)) 

236 [1, 1, 1] 

237 

238 >>> from sympy.abc import a,b,c,d 

239 >>> m1 = Matrix([[a, b], [c, d]]) 

240 >>> tensordiagonal(m1, [0, 1]) 

241 [a, d] 

242 

243 In case of higher dimensional arrays, the diagonalized out dimensions 

244 are appended removed and appended as a single dimension at the end: 

245 

246 >>> A = Array(range(18), (3, 2, 3)) 

247 >>> A 

248 [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]] 

249 >>> tensordiagonal(A, (0, 2)) 

250 [[0, 7, 14], [3, 10, 17]] 

251 >>> from sympy import permutedims 

252 >>> tensordiagonal(A, (0, 2)) == permutedims(Array([A[0, :, 0], A[1, :, 1], A[2, :, 2]]), [1, 0]) 

253 True 

254 

255 See Also 

256 ======== 

257 

258 sympy.tensor.array.expressions.array_expressions.ArrayDiagonal 

259 

260 """ 

261 if any(len(i) <= 1 for i in diagonal_axes): 

262 raise ValueError("need at least two axes to diagonalize") 

263 

264 from sympy.tensor.array.expressions.array_expressions import _ArrayExpr 

265 from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract 

266 from sympy.tensor.array.expressions.array_expressions import ArrayDiagonal, _array_diagonal 

267 from sympy.matrices.expressions.matexpr import MatrixSymbol 

268 if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)): 

269 return _array_diagonal(array, *diagonal_axes) 

270 

271 ArrayDiagonal._validate(array, *diagonal_axes) 

272 

273 array, remaining_indices, remaining_shape, diagonal_deltas = _util_contraction_diagonal(array, *diagonal_axes) 

274 

275 # Compute the diagonalized array: 

276 # 

277 # 1. external for loops on all undiagonalized indices. 

278 # Undiagonalized indices are determined by the combinatorial product of 

279 # the absolute positions of the remaining indices. 

280 # 2. internal loop on all diagonal indices. 

281 # It appends the values of the absolute diagonalized index and the absolute 

282 # undiagonalized index for the external loop. 

283 diagonalized_array = [] 

284 diagonal_shape = [len(i) for i in diagonal_deltas] 

285 for icontrib in itertools.product(*remaining_indices): 

286 index_base_position = sum(icontrib) 

287 isum = [] 

288 for sum_to_index in itertools.product(*diagonal_deltas): 

289 idx = array._get_tuple_index(index_base_position + sum(sum_to_index)) 

290 isum.append(array[idx]) 

291 

292 isum = type(array)(isum).reshape(*diagonal_shape) 

293 diagonalized_array.append(isum) 

294 

295 return type(array)(diagonalized_array, remaining_shape + diagonal_shape) 

296 

297 

298def derive_by_array(expr, dx): 

299 r""" 

300 Derivative by arrays. Supports both arrays and scalars. 

301 

302 The equivalent operator for array expressions is ``array_derive``. 

303 

304 Explanation 

305 =========== 

306 

307 Given the array `A_{i_1, \ldots, i_N}` and the array `X_{j_1, \ldots, j_M}` 

308 this function will return a new array `B` defined by 

309 

310 `B_{j_1,\ldots,j_M,i_1,\ldots,i_N} := \frac{\partial A_{i_1,\ldots,i_N}}{\partial X_{j_1,\ldots,j_M}}` 

311 

312 Examples 

313 ======== 

314 

315 >>> from sympy import derive_by_array 

316 >>> from sympy.abc import x, y, z, t 

317 >>> from sympy import cos 

318 >>> derive_by_array(cos(x*t), x) 

319 -t*sin(t*x) 

320 >>> derive_by_array(cos(x*t), [x, y, z, t]) 

321 [-t*sin(t*x), 0, 0, -x*sin(t*x)] 

322 >>> derive_by_array([x, y**2*z], [[x, y], [z, t]]) 

323 [[[1, 0], [0, 2*y*z]], [[0, y**2], [0, 0]]] 

324 

325 """ 

326 from sympy.matrices import MatrixBase 

327 from sympy.tensor.array import SparseNDimArray 

328 array_types = (Iterable, MatrixBase, NDimArray) 

329 

330 if isinstance(dx, array_types): 

331 dx = ImmutableDenseNDimArray(dx) 

332 for i in dx: 

333 if not i._diff_wrt: 

334 raise ValueError("cannot derive by this array") 

335 

336 if isinstance(expr, array_types): 

337 if isinstance(expr, NDimArray): 

338 expr = expr.as_immutable() 

339 else: 

340 expr = ImmutableDenseNDimArray(expr) 

341 

342 if isinstance(dx, array_types): 

343 if isinstance(expr, SparseNDimArray): 

344 lp = len(expr) 

345 new_array = {k + i*lp: v 

346 for i, x in enumerate(Flatten(dx)) 

347 for k, v in expr.diff(x)._sparse_array.items()} 

348 else: 

349 new_array = [[y.diff(x) for y in Flatten(expr)] for x in Flatten(dx)] 

350 return type(expr)(new_array, dx.shape + expr.shape) 

351 else: 

352 return expr.diff(dx) 

353 else: 

354 expr = _sympify(expr) 

355 if isinstance(dx, array_types): 

356 return ImmutableDenseNDimArray([expr.diff(i) for i in Flatten(dx)], dx.shape) 

357 else: 

358 dx = _sympify(dx) 

359 return diff(expr, dx) 

360 

361 

362def permutedims(expr, perm=None, index_order_old=None, index_order_new=None): 

363 """ 

364 Permutes the indices of an array. 

365 

366 Parameter specifies the permutation of the indices. 

367 

368 The equivalent operator for array expressions is ``PermuteDims``, which can 

369 be used to keep the expression unevaluated. 

370 

371 Examples 

372 ======== 

373 

374 >>> from sympy.abc import x, y, z, t 

375 >>> from sympy import sin 

376 >>> from sympy import Array, permutedims 

377 >>> a = Array([[x, y, z], [t, sin(x), 0]]) 

378 >>> a 

379 [[x, y, z], [t, sin(x), 0]] 

380 >>> permutedims(a, (1, 0)) 

381 [[x, t], [y, sin(x)], [z, 0]] 

382 

383 If the array is of second order, ``transpose`` can be used: 

384 

385 >>> from sympy import transpose 

386 >>> transpose(a) 

387 [[x, t], [y, sin(x)], [z, 0]] 

388 

389 Examples on higher dimensions: 

390 

391 >>> b = Array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) 

392 >>> permutedims(b, (2, 1, 0)) 

393 [[[1, 5], [3, 7]], [[2, 6], [4, 8]]] 

394 >>> permutedims(b, (1, 2, 0)) 

395 [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] 

396 

397 An alternative way to specify the same permutations as in the previous 

398 lines involves passing the *old* and *new* indices, either as a list or as 

399 a string: 

400 

401 >>> permutedims(b, index_order_old="cba", index_order_new="abc") 

402 [[[1, 5], [3, 7]], [[2, 6], [4, 8]]] 

403 >>> permutedims(b, index_order_old="cab", index_order_new="abc") 

404 [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] 

405 

406 ``Permutation`` objects are also allowed: 

407 

408 >>> from sympy.combinatorics import Permutation 

409 >>> permutedims(b, Permutation([1, 2, 0])) 

410 [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] 

411 

412 See Also 

413 ======== 

414 

415 sympy.tensor.array.expressions.array_expressions.PermuteDims 

416 

417 """ 

418 from sympy.tensor.array import SparseNDimArray 

419 

420 from sympy.tensor.array.expressions.array_expressions import _ArrayExpr 

421 from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract 

422 from sympy.tensor.array.expressions.array_expressions import _permute_dims 

423 from sympy.matrices.expressions.matexpr import MatrixSymbol 

424 from sympy.tensor.array.expressions import PermuteDims 

425 from sympy.tensor.array.expressions.array_expressions import get_rank 

426 perm = PermuteDims._get_permutation_from_arguments(perm, index_order_old, index_order_new, get_rank(expr)) 

427 if isinstance(expr, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)): 

428 return _permute_dims(expr, perm) 

429 

430 if not isinstance(expr, NDimArray): 

431 expr = ImmutableDenseNDimArray(expr) 

432 

433 from sympy.combinatorics import Permutation 

434 if not isinstance(perm, Permutation): 

435 perm = Permutation(list(perm)) 

436 

437 if perm.size != expr.rank(): 

438 raise ValueError("wrong permutation size") 

439 

440 # Get the inverse permutation: 

441 iperm = ~perm 

442 new_shape = perm(expr.shape) 

443 

444 if isinstance(expr, SparseNDimArray): 

445 return type(expr)({tuple(perm(expr._get_tuple_index(k))): v 

446 for k, v in expr._sparse_array.items()}, new_shape) 

447 

448 indices_span = perm([range(i) for i in expr.shape]) 

449 

450 new_array = [None]*len(expr) 

451 for i, idx in enumerate(itertools.product(*indices_span)): 

452 t = iperm(idx) 

453 new_array[i] = expr[t] 

454 

455 return type(expr)(new_array, new_shape) 

456 

457 

458class Flatten(Printable): 

459 """ 

460 Flatten an iterable object to a list in a lazy-evaluation way. 

461 

462 Notes 

463 ===== 

464 

465 This class is an iterator with which the memory cost can be economised. 

466 Optimisation has been considered to ameliorate the performance for some 

467 specific data types like DenseNDimArray and SparseNDimArray. 

468 

469 Examples 

470 ======== 

471 

472 >>> from sympy.tensor.array.arrayop import Flatten 

473 >>> from sympy.tensor.array import Array 

474 >>> A = Array(range(6)).reshape(2, 3) 

475 >>> Flatten(A) 

476 Flatten([[0, 1, 2], [3, 4, 5]]) 

477 >>> [i for i in Flatten(A)] 

478 [0, 1, 2, 3, 4, 5] 

479 """ 

480 def __init__(self, iterable): 

481 from sympy.matrices.matrices import MatrixBase 

482 from sympy.tensor.array import NDimArray 

483 

484 if not isinstance(iterable, (Iterable, MatrixBase)): 

485 raise NotImplementedError("Data type not yet supported") 

486 

487 if isinstance(iterable, list): 

488 iterable = NDimArray(iterable) 

489 

490 self._iter = iterable 

491 self._idx = 0 

492 

493 def __iter__(self): 

494 return self 

495 

496 def __next__(self): 

497 from sympy.matrices.matrices import MatrixBase 

498 

499 if len(self._iter) > self._idx: 

500 if isinstance(self._iter, DenseNDimArray): 

501 result = self._iter._array[self._idx] 

502 

503 elif isinstance(self._iter, SparseNDimArray): 

504 if self._idx in self._iter._sparse_array: 

505 result = self._iter._sparse_array[self._idx] 

506 else: 

507 result = 0 

508 

509 elif isinstance(self._iter, MatrixBase): 

510 result = self._iter[self._idx] 

511 

512 elif hasattr(self._iter, '__next__'): 

513 result = next(self._iter) 

514 

515 else: 

516 result = self._iter[self._idx] 

517 

518 else: 

519 raise StopIteration 

520 

521 self._idx += 1 

522 return result 

523 

524 def next(self): 

525 return self.__next__() 

526 

527 def _sympystr(self, printer): 

528 return type(self).__name__ + '(' + printer._print(self._iter) + ')'