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

290 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1from sympy.core.basic import Basic 

2from sympy.core.containers import (Dict, Tuple) 

3from sympy.core.expr import Expr 

4from sympy.core.kind import Kind, NumberKind, UndefinedKind 

5from sympy.core.numbers import Integer 

6from sympy.core.singleton import S 

7from sympy.core.sympify import sympify 

8from sympy.external.gmpy import SYMPY_INTS 

9from sympy.printing.defaults import Printable 

10 

11import itertools 

12from collections.abc import Iterable 

13 

14 

15class ArrayKind(Kind): 

16 """ 

17 Kind for N-dimensional array in SymPy. 

18 

19 This kind represents the multidimensional array that algebraic 

20 operations are defined. Basic class for this kind is ``NDimArray``, 

21 but any expression representing the array can have this. 

22 

23 Parameters 

24 ========== 

25 

26 element_kind : Kind 

27 Kind of the element. Default is :obj:NumberKind `<sympy.core.kind.NumberKind>`, 

28 which means that the array contains only numbers. 

29 

30 Examples 

31 ======== 

32 

33 Any instance of array class has ``ArrayKind``. 

34 

35 >>> from sympy import NDimArray 

36 >>> NDimArray([1,2,3]).kind 

37 ArrayKind(NumberKind) 

38 

39 Although expressions representing an array may be not instance of 

40 array class, it will have ``ArrayKind`` as well. 

41 

42 >>> from sympy import Integral 

43 >>> from sympy.tensor.array import NDimArray 

44 >>> from sympy.abc import x 

45 >>> intA = Integral(NDimArray([1,2,3]), x) 

46 >>> isinstance(intA, NDimArray) 

47 False 

48 >>> intA.kind 

49 ArrayKind(NumberKind) 

50 

51 Use ``isinstance()`` to check for ``ArrayKind` without specifying 

52 the element kind. Use ``is`` with specifying the element kind. 

53 

54 >>> from sympy.tensor.array import ArrayKind 

55 >>> from sympy.core import NumberKind 

56 >>> boolA = NDimArray([True, False]) 

57 >>> isinstance(boolA.kind, ArrayKind) 

58 True 

59 >>> boolA.kind is ArrayKind(NumberKind) 

60 False 

61 

62 See Also 

63 ======== 

64 

65 shape : Function to return the shape of objects with ``MatrixKind``. 

66 

67 """ 

68 def __new__(cls, element_kind=NumberKind): 

69 obj = super().__new__(cls, element_kind) 

70 obj.element_kind = element_kind 

71 return obj 

72 

73 def __repr__(self): 

74 return "ArrayKind(%s)" % self.element_kind 

75 

76 @classmethod 

77 def _union(cls, kinds) -> 'ArrayKind': 

78 elem_kinds = {e.kind for e in kinds} 

79 if len(elem_kinds) == 1: 

80 elemkind, = elem_kinds 

81 else: 

82 elemkind = UndefinedKind 

83 return ArrayKind(elemkind) 

84 

85 

86class NDimArray(Printable): 

87 """N-dimensional array. 

88 

89 Examples 

90 ======== 

91 

92 Create an N-dim array of zeros: 

93 

94 >>> from sympy import MutableDenseNDimArray 

95 >>> a = MutableDenseNDimArray.zeros(2, 3, 4) 

96 >>> a 

97 [[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]] 

98 

99 Create an N-dim array from a list; 

100 

101 >>> a = MutableDenseNDimArray([[2, 3], [4, 5]]) 

102 >>> a 

103 [[2, 3], [4, 5]] 

104 

105 >>> b = MutableDenseNDimArray([[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]) 

106 >>> b 

107 [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]] 

108 

109 Create an N-dim array from a flat list with dimension shape: 

110 

111 >>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3)) 

112 >>> a 

113 [[1, 2, 3], [4, 5, 6]] 

114 

115 Create an N-dim array from a matrix: 

116 

117 >>> from sympy import Matrix 

118 >>> a = Matrix([[1,2],[3,4]]) 

119 >>> a 

120 Matrix([ 

121 [1, 2], 

122 [3, 4]]) 

123 >>> b = MutableDenseNDimArray(a) 

124 >>> b 

125 [[1, 2], [3, 4]] 

126 

127 Arithmetic operations on N-dim arrays 

128 

129 >>> a = MutableDenseNDimArray([1, 1, 1, 1], (2, 2)) 

130 >>> b = MutableDenseNDimArray([4, 4, 4, 4], (2, 2)) 

131 >>> c = a + b 

132 >>> c 

133 [[5, 5], [5, 5]] 

134 >>> a - b 

135 [[-3, -3], [-3, -3]] 

136 

137 """ 

138 

139 _diff_wrt = True 

140 is_scalar = False 

141 

142 def __new__(cls, iterable, shape=None, **kwargs): 

143 from sympy.tensor.array import ImmutableDenseNDimArray 

144 return ImmutableDenseNDimArray(iterable, shape, **kwargs) 

145 

146 def __getitem__(self, index): 

147 raise NotImplementedError("A subclass of NDimArray should implement __getitem__") 

148 

149 def _parse_index(self, index): 

150 if isinstance(index, (SYMPY_INTS, Integer)): 

151 if index >= self._loop_size: 

152 raise ValueError("Only a tuple index is accepted") 

153 return index 

154 

155 if self._loop_size == 0: 

156 raise ValueError("Index not valid with an empty array") 

157 

158 if len(index) != self._rank: 

159 raise ValueError('Wrong number of array axes') 

160 

161 real_index = 0 

162 # check if input index can exist in current indexing 

163 for i in range(self._rank): 

164 if (index[i] >= self.shape[i]) or (index[i] < -self.shape[i]): 

165 raise ValueError('Index ' + str(index) + ' out of border') 

166 if index[i] < 0: 

167 real_index += 1 

168 real_index = real_index*self.shape[i] + index[i] 

169 

170 return real_index 

171 

172 def _get_tuple_index(self, integer_index): 

173 index = [] 

174 for i, sh in enumerate(reversed(self.shape)): 

175 index.append(integer_index % sh) 

176 integer_index //= sh 

177 index.reverse() 

178 return tuple(index) 

179 

180 def _check_symbolic_index(self, index): 

181 # Check if any index is symbolic: 

182 tuple_index = (index if isinstance(index, tuple) else (index,)) 

183 if any((isinstance(i, Expr) and (not i.is_number)) for i in tuple_index): 

184 for i, nth_dim in zip(tuple_index, self.shape): 

185 if ((i < 0) == True) or ((i >= nth_dim) == True): 

186 raise ValueError("index out of range") 

187 from sympy.tensor import Indexed 

188 return Indexed(self, *tuple_index) 

189 return None 

190 

191 def _setter_iterable_check(self, value): 

192 from sympy.matrices.matrices import MatrixBase 

193 if isinstance(value, (Iterable, MatrixBase, NDimArray)): 

194 raise NotImplementedError 

195 

196 @classmethod 

197 def _scan_iterable_shape(cls, iterable): 

198 def f(pointer): 

199 if not isinstance(pointer, Iterable): 

200 return [pointer], () 

201 

202 if len(pointer) == 0: 

203 return [], (0,) 

204 

205 result = [] 

206 elems, shapes = zip(*[f(i) for i in pointer]) 

207 if len(set(shapes)) != 1: 

208 raise ValueError("could not determine shape unambiguously") 

209 for i in elems: 

210 result.extend(i) 

211 return result, (len(shapes),)+shapes[0] 

212 

213 return f(iterable) 

214 

215 @classmethod 

216 def _handle_ndarray_creation_inputs(cls, iterable=None, shape=None, **kwargs): 

217 from sympy.matrices.matrices import MatrixBase 

218 from sympy.tensor.array import SparseNDimArray 

219 

220 if shape is None: 

221 if iterable is None: 

222 shape = () 

223 iterable = () 

224 # Construction of a sparse array from a sparse array 

225 elif isinstance(iterable, SparseNDimArray): 

226 return iterable._shape, iterable._sparse_array 

227 

228 # Construct N-dim array from another N-dim array: 

229 elif isinstance(iterable, NDimArray): 

230 shape = iterable.shape 

231 

232 # Construct N-dim array from an iterable (numpy arrays included): 

233 elif isinstance(iterable, Iterable): 

234 iterable, shape = cls._scan_iterable_shape(iterable) 

235 

236 # Construct N-dim array from a Matrix: 

237 elif isinstance(iterable, MatrixBase): 

238 shape = iterable.shape 

239 

240 else: 

241 shape = () 

242 iterable = (iterable,) 

243 

244 if isinstance(iterable, (Dict, dict)) and shape is not None: 

245 new_dict = iterable.copy() 

246 for k, v in new_dict.items(): 

247 if isinstance(k, (tuple, Tuple)): 

248 new_key = 0 

249 for i, idx in enumerate(k): 

250 new_key = new_key * shape[i] + idx 

251 iterable[new_key] = iterable[k] 

252 del iterable[k] 

253 

254 if isinstance(shape, (SYMPY_INTS, Integer)): 

255 shape = (shape,) 

256 

257 if not all(isinstance(dim, (SYMPY_INTS, Integer)) for dim in shape): 

258 raise TypeError("Shape should contain integers only.") 

259 

260 return tuple(shape), iterable 

261 

262 def __len__(self): 

263 """Overload common function len(). Returns number of elements in array. 

264 

265 Examples 

266 ======== 

267 

268 >>> from sympy import MutableDenseNDimArray 

269 >>> a = MutableDenseNDimArray.zeros(3, 3) 

270 >>> a 

271 [[0, 0, 0], [0, 0, 0], [0, 0, 0]] 

272 >>> len(a) 

273 9 

274 

275 """ 

276 return self._loop_size 

277 

278 @property 

279 def shape(self): 

280 """ 

281 Returns array shape (dimension). 

282 

283 Examples 

284 ======== 

285 

286 >>> from sympy import MutableDenseNDimArray 

287 >>> a = MutableDenseNDimArray.zeros(3, 3) 

288 >>> a.shape 

289 (3, 3) 

290 

291 """ 

292 return self._shape 

293 

294 def rank(self): 

295 """ 

296 Returns rank of array. 

297 

298 Examples 

299 ======== 

300 

301 >>> from sympy import MutableDenseNDimArray 

302 >>> a = MutableDenseNDimArray.zeros(3,4,5,6,3) 

303 >>> a.rank() 

304 5 

305 

306 """ 

307 return self._rank 

308 

309 def diff(self, *args, **kwargs): 

310 """ 

311 Calculate the derivative of each element in the array. 

312 

313 Examples 

314 ======== 

315 

316 >>> from sympy import ImmutableDenseNDimArray 

317 >>> from sympy.abc import x, y 

318 >>> M = ImmutableDenseNDimArray([[x, y], [1, x*y]]) 

319 >>> M.diff(x) 

320 [[1, 0], [0, y]] 

321 

322 """ 

323 from sympy.tensor.array.array_derivatives import ArrayDerivative 

324 kwargs.setdefault('evaluate', True) 

325 return ArrayDerivative(self.as_immutable(), *args, **kwargs) 

326 

327 def _eval_derivative(self, base): 

328 # Types are (base: scalar, self: array) 

329 return self.applyfunc(lambda x: base.diff(x)) 

330 

331 def _eval_derivative_n_times(self, s, n): 

332 return Basic._eval_derivative_n_times(self, s, n) 

333 

334 def applyfunc(self, f): 

335 """Apply a function to each element of the N-dim array. 

336 

337 Examples 

338 ======== 

339 

340 >>> from sympy import ImmutableDenseNDimArray 

341 >>> m = ImmutableDenseNDimArray([i*2+j for i in range(2) for j in range(2)], (2, 2)) 

342 >>> m 

343 [[0, 1], [2, 3]] 

344 >>> m.applyfunc(lambda i: 2*i) 

345 [[0, 2], [4, 6]] 

346 """ 

347 from sympy.tensor.array import SparseNDimArray 

348 from sympy.tensor.array.arrayop import Flatten 

349 

350 if isinstance(self, SparseNDimArray) and f(S.Zero) == 0: 

351 return type(self)({k: f(v) for k, v in self._sparse_array.items() if f(v) != 0}, self.shape) 

352 

353 return type(self)(map(f, Flatten(self)), self.shape) 

354 

355 def _sympystr(self, printer): 

356 def f(sh, shape_left, i, j): 

357 if len(shape_left) == 1: 

358 return "["+", ".join([printer._print(self[self._get_tuple_index(e)]) for e in range(i, j)])+"]" 

359 

360 sh //= shape_left[0] 

361 return "[" + ", ".join([f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh) for e in range(shape_left[0])]) + "]" # + "\n"*len(shape_left) 

362 

363 if self.rank() == 0: 

364 return printer._print(self[()]) 

365 

366 return f(self._loop_size, self.shape, 0, self._loop_size) 

367 

368 def tolist(self): 

369 """ 

370 Converting MutableDenseNDimArray to one-dim list 

371 

372 Examples 

373 ======== 

374 

375 >>> from sympy import MutableDenseNDimArray 

376 >>> a = MutableDenseNDimArray([1, 2, 3, 4], (2, 2)) 

377 >>> a 

378 [[1, 2], [3, 4]] 

379 >>> b = a.tolist() 

380 >>> b 

381 [[1, 2], [3, 4]] 

382 """ 

383 

384 def f(sh, shape_left, i, j): 

385 if len(shape_left) == 1: 

386 return [self[self._get_tuple_index(e)] for e in range(i, j)] 

387 result = [] 

388 sh //= shape_left[0] 

389 for e in range(shape_left[0]): 

390 result.append(f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh)) 

391 return result 

392 

393 return f(self._loop_size, self.shape, 0, self._loop_size) 

394 

395 def __add__(self, other): 

396 from sympy.tensor.array.arrayop import Flatten 

397 

398 if not isinstance(other, NDimArray): 

399 return NotImplemented 

400 

401 if self.shape != other.shape: 

402 raise ValueError("array shape mismatch") 

403 result_list = [i+j for i,j in zip(Flatten(self), Flatten(other))] 

404 

405 return type(self)(result_list, self.shape) 

406 

407 def __sub__(self, other): 

408 from sympy.tensor.array.arrayop import Flatten 

409 

410 if not isinstance(other, NDimArray): 

411 return NotImplemented 

412 

413 if self.shape != other.shape: 

414 raise ValueError("array shape mismatch") 

415 result_list = [i-j for i,j in zip(Flatten(self), Flatten(other))] 

416 

417 return type(self)(result_list, self.shape) 

418 

419 def __mul__(self, other): 

420 from sympy.matrices.matrices import MatrixBase 

421 from sympy.tensor.array import SparseNDimArray 

422 from sympy.tensor.array.arrayop import Flatten 

423 

424 if isinstance(other, (Iterable, NDimArray, MatrixBase)): 

425 raise ValueError("scalar expected, use tensorproduct(...) for tensorial product") 

426 

427 other = sympify(other) 

428 if isinstance(self, SparseNDimArray): 

429 if other.is_zero: 

430 return type(self)({}, self.shape) 

431 return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape) 

432 

433 result_list = [i*other for i in Flatten(self)] 

434 return type(self)(result_list, self.shape) 

435 

436 def __rmul__(self, other): 

437 from sympy.matrices.matrices import MatrixBase 

438 from sympy.tensor.array import SparseNDimArray 

439 from sympy.tensor.array.arrayop import Flatten 

440 

441 if isinstance(other, (Iterable, NDimArray, MatrixBase)): 

442 raise ValueError("scalar expected, use tensorproduct(...) for tensorial product") 

443 

444 other = sympify(other) 

445 if isinstance(self, SparseNDimArray): 

446 if other.is_zero: 

447 return type(self)({}, self.shape) 

448 return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape) 

449 

450 result_list = [other*i for i in Flatten(self)] 

451 return type(self)(result_list, self.shape) 

452 

453 def __truediv__(self, other): 

454 from sympy.matrices.matrices import MatrixBase 

455 from sympy.tensor.array import SparseNDimArray 

456 from sympy.tensor.array.arrayop import Flatten 

457 

458 if isinstance(other, (Iterable, NDimArray, MatrixBase)): 

459 raise ValueError("scalar expected") 

460 

461 other = sympify(other) 

462 if isinstance(self, SparseNDimArray) and other != S.Zero: 

463 return type(self)({k: v/other for (k, v) in self._sparse_array.items()}, self.shape) 

464 

465 result_list = [i/other for i in Flatten(self)] 

466 return type(self)(result_list, self.shape) 

467 

468 def __rtruediv__(self, other): 

469 raise NotImplementedError('unsupported operation on NDimArray') 

470 

471 def __neg__(self): 

472 from sympy.tensor.array import SparseNDimArray 

473 from sympy.tensor.array.arrayop import Flatten 

474 

475 if isinstance(self, SparseNDimArray): 

476 return type(self)({k: -v for (k, v) in self._sparse_array.items()}, self.shape) 

477 

478 result_list = [-i for i in Flatten(self)] 

479 return type(self)(result_list, self.shape) 

480 

481 def __iter__(self): 

482 def iterator(): 

483 if self._shape: 

484 for i in range(self._shape[0]): 

485 yield self[i] 

486 else: 

487 yield self[()] 

488 

489 return iterator() 

490 

491 def __eq__(self, other): 

492 """ 

493 NDimArray instances can be compared to each other. 

494 Instances equal if they have same shape and data. 

495 

496 Examples 

497 ======== 

498 

499 >>> from sympy import MutableDenseNDimArray 

500 >>> a = MutableDenseNDimArray.zeros(2, 3) 

501 >>> b = MutableDenseNDimArray.zeros(2, 3) 

502 >>> a == b 

503 True 

504 >>> c = a.reshape(3, 2) 

505 >>> c == b 

506 False 

507 >>> a[0,0] = 1 

508 >>> b[0,0] = 2 

509 >>> a == b 

510 False 

511 """ 

512 from sympy.tensor.array import SparseNDimArray 

513 if not isinstance(other, NDimArray): 

514 return False 

515 

516 if not self.shape == other.shape: 

517 return False 

518 

519 if isinstance(self, SparseNDimArray) and isinstance(other, SparseNDimArray): 

520 return dict(self._sparse_array) == dict(other._sparse_array) 

521 

522 return list(self) == list(other) 

523 

524 def __ne__(self, other): 

525 return not self == other 

526 

527 def _eval_transpose(self): 

528 if self.rank() != 2: 

529 raise ValueError("array rank not 2") 

530 from .arrayop import permutedims 

531 return permutedims(self, (1, 0)) 

532 

533 def transpose(self): 

534 return self._eval_transpose() 

535 

536 def _eval_conjugate(self): 

537 from sympy.tensor.array.arrayop import Flatten 

538 

539 return self.func([i.conjugate() for i in Flatten(self)], self.shape) 

540 

541 def conjugate(self): 

542 return self._eval_conjugate() 

543 

544 def _eval_adjoint(self): 

545 return self.transpose().conjugate() 

546 

547 def adjoint(self): 

548 return self._eval_adjoint() 

549 

550 def _slice_expand(self, s, dim): 

551 if not isinstance(s, slice): 

552 return (s,) 

553 start, stop, step = s.indices(dim) 

554 return [start + i*step for i in range((stop-start)//step)] 

555 

556 def _get_slice_data_for_array_access(self, index): 

557 sl_factors = [self._slice_expand(i, dim) for (i, dim) in zip(index, self.shape)] 

558 eindices = itertools.product(*sl_factors) 

559 return sl_factors, eindices 

560 

561 def _get_slice_data_for_array_assignment(self, index, value): 

562 if not isinstance(value, NDimArray): 

563 value = type(self)(value) 

564 sl_factors, eindices = self._get_slice_data_for_array_access(index) 

565 slice_offsets = [min(i) if isinstance(i, list) else None for i in sl_factors] 

566 # TODO: add checks for dimensions for `value`? 

567 return value, eindices, slice_offsets 

568 

569 @classmethod 

570 def _check_special_bounds(cls, flat_list, shape): 

571 if shape == () and len(flat_list) != 1: 

572 raise ValueError("arrays without shape need one scalar value") 

573 if shape == (0,) and len(flat_list) > 0: 

574 raise ValueError("if array shape is (0,) there cannot be elements") 

575 

576 def _check_index_for_getitem(self, index): 

577 if isinstance(index, (SYMPY_INTS, Integer, slice)): 

578 index = (index,) 

579 

580 if len(index) < self.rank(): 

581 index = tuple(index) + \ 

582 tuple(slice(None) for i in range(len(index), self.rank())) 

583 

584 if len(index) > self.rank(): 

585 raise ValueError('Dimension of index greater than rank of array') 

586 

587 return index 

588 

589 

590class ImmutableNDimArray(NDimArray, Basic): 

591 _op_priority = 11.0 

592 

593 def __hash__(self): 

594 return Basic.__hash__(self) 

595 

596 def as_immutable(self): 

597 return self 

598 

599 def as_mutable(self): 

600 raise NotImplementedError("abstract method")