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
« 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
11import itertools
12from collections.abc import Iterable
15class ArrayKind(Kind):
16 """
17 Kind for N-dimensional array in SymPy.
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.
23 Parameters
24 ==========
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.
30 Examples
31 ========
33 Any instance of array class has ``ArrayKind``.
35 >>> from sympy import NDimArray
36 >>> NDimArray([1,2,3]).kind
37 ArrayKind(NumberKind)
39 Although expressions representing an array may be not instance of
40 array class, it will have ``ArrayKind`` as well.
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)
51 Use ``isinstance()`` to check for ``ArrayKind` without specifying
52 the element kind. Use ``is`` with specifying the element kind.
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
62 See Also
63 ========
65 shape : Function to return the shape of objects with ``MatrixKind``.
67 """
68 def __new__(cls, element_kind=NumberKind):
69 obj = super().__new__(cls, element_kind)
70 obj.element_kind = element_kind
71 return obj
73 def __repr__(self):
74 return "ArrayKind(%s)" % self.element_kind
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)
86class NDimArray(Printable):
87 """N-dimensional array.
89 Examples
90 ========
92 Create an N-dim array of zeros:
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]]]
99 Create an N-dim array from a list;
101 >>> a = MutableDenseNDimArray([[2, 3], [4, 5]])
102 >>> a
103 [[2, 3], [4, 5]]
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]]]
109 Create an N-dim array from a flat list with dimension shape:
111 >>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3))
112 >>> a
113 [[1, 2, 3], [4, 5, 6]]
115 Create an N-dim array from a matrix:
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]]
127 Arithmetic operations on N-dim arrays
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]]
137 """
139 _diff_wrt = True
140 is_scalar = False
142 def __new__(cls, iterable, shape=None, **kwargs):
143 from sympy.tensor.array import ImmutableDenseNDimArray
144 return ImmutableDenseNDimArray(iterable, shape, **kwargs)
146 def __getitem__(self, index):
147 raise NotImplementedError("A subclass of NDimArray should implement __getitem__")
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
155 if self._loop_size == 0:
156 raise ValueError("Index not valid with an empty array")
158 if len(index) != self._rank:
159 raise ValueError('Wrong number of array axes')
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]
170 return real_index
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)
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
191 def _setter_iterable_check(self, value):
192 from sympy.matrices.matrices import MatrixBase
193 if isinstance(value, (Iterable, MatrixBase, NDimArray)):
194 raise NotImplementedError
196 @classmethod
197 def _scan_iterable_shape(cls, iterable):
198 def f(pointer):
199 if not isinstance(pointer, Iterable):
200 return [pointer], ()
202 if len(pointer) == 0:
203 return [], (0,)
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]
213 return f(iterable)
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
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
228 # Construct N-dim array from another N-dim array:
229 elif isinstance(iterable, NDimArray):
230 shape = iterable.shape
232 # Construct N-dim array from an iterable (numpy arrays included):
233 elif isinstance(iterable, Iterable):
234 iterable, shape = cls._scan_iterable_shape(iterable)
236 # Construct N-dim array from a Matrix:
237 elif isinstance(iterable, MatrixBase):
238 shape = iterable.shape
240 else:
241 shape = ()
242 iterable = (iterable,)
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]
254 if isinstance(shape, (SYMPY_INTS, Integer)):
255 shape = (shape,)
257 if not all(isinstance(dim, (SYMPY_INTS, Integer)) for dim in shape):
258 raise TypeError("Shape should contain integers only.")
260 return tuple(shape), iterable
262 def __len__(self):
263 """Overload common function len(). Returns number of elements in array.
265 Examples
266 ========
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
275 """
276 return self._loop_size
278 @property
279 def shape(self):
280 """
281 Returns array shape (dimension).
283 Examples
284 ========
286 >>> from sympy import MutableDenseNDimArray
287 >>> a = MutableDenseNDimArray.zeros(3, 3)
288 >>> a.shape
289 (3, 3)
291 """
292 return self._shape
294 def rank(self):
295 """
296 Returns rank of array.
298 Examples
299 ========
301 >>> from sympy import MutableDenseNDimArray
302 >>> a = MutableDenseNDimArray.zeros(3,4,5,6,3)
303 >>> a.rank()
304 5
306 """
307 return self._rank
309 def diff(self, *args, **kwargs):
310 """
311 Calculate the derivative of each element in the array.
313 Examples
314 ========
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]]
322 """
323 from sympy.tensor.array.array_derivatives import ArrayDerivative
324 kwargs.setdefault('evaluate', True)
325 return ArrayDerivative(self.as_immutable(), *args, **kwargs)
327 def _eval_derivative(self, base):
328 # Types are (base: scalar, self: array)
329 return self.applyfunc(lambda x: base.diff(x))
331 def _eval_derivative_n_times(self, s, n):
332 return Basic._eval_derivative_n_times(self, s, n)
334 def applyfunc(self, f):
335 """Apply a function to each element of the N-dim array.
337 Examples
338 ========
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
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)
353 return type(self)(map(f, Flatten(self)), self.shape)
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)])+"]"
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)
363 if self.rank() == 0:
364 return printer._print(self[()])
366 return f(self._loop_size, self.shape, 0, self._loop_size)
368 def tolist(self):
369 """
370 Converting MutableDenseNDimArray to one-dim list
372 Examples
373 ========
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 """
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
393 return f(self._loop_size, self.shape, 0, self._loop_size)
395 def __add__(self, other):
396 from sympy.tensor.array.arrayop import Flatten
398 if not isinstance(other, NDimArray):
399 return NotImplemented
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))]
405 return type(self)(result_list, self.shape)
407 def __sub__(self, other):
408 from sympy.tensor.array.arrayop import Flatten
410 if not isinstance(other, NDimArray):
411 return NotImplemented
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))]
417 return type(self)(result_list, self.shape)
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
424 if isinstance(other, (Iterable, NDimArray, MatrixBase)):
425 raise ValueError("scalar expected, use tensorproduct(...) for tensorial product")
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)
433 result_list = [i*other for i in Flatten(self)]
434 return type(self)(result_list, self.shape)
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
441 if isinstance(other, (Iterable, NDimArray, MatrixBase)):
442 raise ValueError("scalar expected, use tensorproduct(...) for tensorial product")
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)
450 result_list = [other*i for i in Flatten(self)]
451 return type(self)(result_list, self.shape)
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
458 if isinstance(other, (Iterable, NDimArray, MatrixBase)):
459 raise ValueError("scalar expected")
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)
465 result_list = [i/other for i in Flatten(self)]
466 return type(self)(result_list, self.shape)
468 def __rtruediv__(self, other):
469 raise NotImplementedError('unsupported operation on NDimArray')
471 def __neg__(self):
472 from sympy.tensor.array import SparseNDimArray
473 from sympy.tensor.array.arrayop import Flatten
475 if isinstance(self, SparseNDimArray):
476 return type(self)({k: -v for (k, v) in self._sparse_array.items()}, self.shape)
478 result_list = [-i for i in Flatten(self)]
479 return type(self)(result_list, self.shape)
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[()]
489 return iterator()
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.
496 Examples
497 ========
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
516 if not self.shape == other.shape:
517 return False
519 if isinstance(self, SparseNDimArray) and isinstance(other, SparseNDimArray):
520 return dict(self._sparse_array) == dict(other._sparse_array)
522 return list(self) == list(other)
524 def __ne__(self, other):
525 return not self == other
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))
533 def transpose(self):
534 return self._eval_transpose()
536 def _eval_conjugate(self):
537 from sympy.tensor.array.arrayop import Flatten
539 return self.func([i.conjugate() for i in Flatten(self)], self.shape)
541 def conjugate(self):
542 return self._eval_conjugate()
544 def _eval_adjoint(self):
545 return self.transpose().conjugate()
547 def adjoint(self):
548 return self._eval_adjoint()
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)]
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
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
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")
576 def _check_index_for_getitem(self, index):
577 if isinstance(index, (SYMPY_INTS, Integer, slice)):
578 index = (index,)
580 if len(index) < self.rank():
581 index = tuple(index) + \
582 tuple(slice(None) for i in range(len(index), self.rank()))
584 if len(index) > self.rank():
585 raise ValueError('Dimension of index greater than rank of array')
587 return index
590class ImmutableNDimArray(NDimArray, Basic):
591 _op_priority = 11.0
593 def __hash__(self):
594 return Basic.__hash__(self)
596 def as_immutable(self):
597 return self
599 def as_mutable(self):
600 raise NotImplementedError("abstract method")