Coverage for /usr/lib/python3/dist-packages/sympy/tensor/indexed.py: 33%
269 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1r"""Module that defines indexed objects.
3The classes ``IndexedBase``, ``Indexed``, and ``Idx`` represent a
4matrix element ``M[i, j]`` as in the following diagram::
6 1) The Indexed class represents the entire indexed object.
7 |
8 ___|___
9 ' '
10 M[i, j]
11 / \__\______
12 | |
13 | |
14 | 2) The Idx class represents indices; each Idx can
15 | optionally contain information about its range.
16 |
17 3) IndexedBase represents the 'stem' of an indexed object, here `M`.
18 The stem used by itself is usually taken to represent the entire
19 array.
21There can be any number of indices on an Indexed object. No
22transformation properties are implemented in these Base objects, but
23implicit contraction of repeated indices is supported.
25Note that the support for complicated (i.e. non-atomic) integer
26expressions as indices is limited. (This should be improved in
27future releases.)
29Examples
30========
32To express the above matrix element example you would write:
34>>> from sympy import symbols, IndexedBase, Idx
35>>> M = IndexedBase('M')
36>>> i, j = symbols('i j', cls=Idx)
37>>> M[i, j]
38M[i, j]
40Repeated indices in a product implies a summation, so to express a
41matrix-vector product in terms of Indexed objects:
43>>> x = IndexedBase('x')
44>>> M[i, j]*x[j]
45M[i, j]*x[j]
47If the indexed objects will be converted to component based arrays, e.g.
48with the code printers or the autowrap framework, you also need to provide
49(symbolic or numerical) dimensions. This can be done by passing an
50optional shape parameter to IndexedBase upon construction:
52>>> dim1, dim2 = symbols('dim1 dim2', integer=True)
53>>> A = IndexedBase('A', shape=(dim1, 2*dim1, dim2))
54>>> A.shape
55(dim1, 2*dim1, dim2)
56>>> A[i, j, 3].shape
57(dim1, 2*dim1, dim2)
59If an IndexedBase object has no shape information, it is assumed that the
60array is as large as the ranges of its indices:
62>>> n, m = symbols('n m', integer=True)
63>>> i = Idx('i', m)
64>>> j = Idx('j', n)
65>>> M[i, j].shape
66(m, n)
67>>> M[i, j].ranges
68[(0, m - 1), (0, n - 1)]
70The above can be compared with the following:
72>>> A[i, 2, j].shape
73(dim1, 2*dim1, dim2)
74>>> A[i, 2, j].ranges
75[(0, m - 1), None, (0, n - 1)]
77To analyze the structure of indexed expressions, you can use the methods
78get_indices() and get_contraction_structure():
80>>> from sympy.tensor import get_indices, get_contraction_structure
81>>> get_indices(A[i, j, j])
82({i}, {})
83>>> get_contraction_structure(A[i, j, j])
84{(j,): {A[i, j, j]}}
86See the appropriate docstrings for a detailed explanation of the output.
87"""
89# TODO: (some ideas for improvement)
90#
91# o test and guarantee numpy compatibility
92# - implement full support for broadcasting
93# - strided arrays
94#
95# o more functions to analyze indexed expressions
96# - identify standard constructs, e.g matrix-vector product in a subexpression
97#
98# o functions to generate component based arrays (numpy and sympy.Matrix)
99# - generate a single array directly from Indexed
100# - convert simple sub-expressions
101#
102# o sophisticated indexing (possibly in subclasses to preserve simplicity)
103# - Idx with range smaller than dimension of Indexed
104# - Idx with stepsize != 1
105# - Idx with step determined by function call
106from collections.abc import Iterable
108from sympy.core.numbers import Number
109from sympy.core.assumptions import StdFactKB
110from sympy.core import Expr, Tuple, sympify, S
111from sympy.core.symbol import _filter_assumptions, Symbol
112from sympy.core.logic import fuzzy_bool, fuzzy_not
113from sympy.core.sympify import _sympify
114from sympy.functions.special.tensor_functions import KroneckerDelta
115from sympy.multipledispatch import dispatch
116from sympy.utilities.iterables import is_sequence, NotIterable
117from sympy.utilities.misc import filldedent
120class IndexException(Exception):
121 pass
124class Indexed(Expr):
125 """Represents a mathematical object with indices.
127 >>> from sympy import Indexed, IndexedBase, Idx, symbols
128 >>> i, j = symbols('i j', cls=Idx)
129 >>> Indexed('A', i, j)
130 A[i, j]
132 It is recommended that ``Indexed`` objects be created by indexing ``IndexedBase``:
133 ``IndexedBase('A')[i, j]`` instead of ``Indexed(IndexedBase('A'), i, j)``.
135 >>> A = IndexedBase('A')
136 >>> a_ij = A[i, j] # Prefer this,
137 >>> b_ij = Indexed(A, i, j) # over this.
138 >>> a_ij == b_ij
139 True
141 """
142 is_commutative = True
143 is_Indexed = True
144 is_symbol = True
145 is_Atom = True
147 def __new__(cls, base, *args, **kw_args):
148 from sympy.tensor.array.ndim_array import NDimArray
149 from sympy.matrices.matrices import MatrixBase
151 if not args:
152 raise IndexException("Indexed needs at least one index.")
153 if isinstance(base, (str, Symbol)):
154 base = IndexedBase(base)
155 elif not hasattr(base, '__getitem__') and not isinstance(base, IndexedBase):
156 raise TypeError(filldedent("""
157 The base can only be replaced with a string, Symbol,
158 IndexedBase or an object with a method for getting
159 items (i.e. an object with a `__getitem__` method).
160 """))
161 args = list(map(sympify, args))
162 if isinstance(base, (NDimArray, Iterable, Tuple, MatrixBase)) and all(i.is_number for i in args):
163 if len(args) == 1:
164 return base[args[0]]
165 else:
166 return base[args]
168 base = _sympify(base)
170 obj = Expr.__new__(cls, base, *args, **kw_args)
172 try:
173 IndexedBase._set_assumptions(obj, base.assumptions0)
174 except AttributeError:
175 IndexedBase._set_assumptions(obj, {})
176 return obj
178 def _hashable_content(self):
179 return super()._hashable_content() + tuple(sorted(self.assumptions0.items()))
181 @property
182 def name(self):
183 return str(self)
185 @property
186 def _diff_wrt(self):
187 """Allow derivatives with respect to an ``Indexed`` object."""
188 return True
190 def _eval_derivative(self, wrt):
191 from sympy.tensor.array.ndim_array import NDimArray
193 if isinstance(wrt, Indexed) and wrt.base == self.base:
194 if len(self.indices) != len(wrt.indices):
195 msg = "Different # of indices: d({!s})/d({!s})".format(self,
196 wrt)
197 raise IndexException(msg)
198 result = S.One
199 for index1, index2 in zip(self.indices, wrt.indices):
200 result *= KroneckerDelta(index1, index2)
201 return result
202 elif isinstance(self.base, NDimArray):
203 from sympy.tensor.array import derive_by_array
204 return Indexed(derive_by_array(self.base, wrt), *self.args[1:])
205 else:
206 if Tuple(self.indices).has(wrt):
207 return S.NaN
208 return S.Zero
210 @property
211 def assumptions0(self):
212 return {k: v for k, v in self._assumptions.items() if v is not None}
214 @property
215 def base(self):
216 """Returns the ``IndexedBase`` of the ``Indexed`` object.
218 Examples
219 ========
221 >>> from sympy import Indexed, IndexedBase, Idx, symbols
222 >>> i, j = symbols('i j', cls=Idx)
223 >>> Indexed('A', i, j).base
224 A
225 >>> B = IndexedBase('B')
226 >>> B == B[i, j].base
227 True
229 """
230 return self.args[0]
232 @property
233 def indices(self):
234 """
235 Returns the indices of the ``Indexed`` object.
237 Examples
238 ========
240 >>> from sympy import Indexed, Idx, symbols
241 >>> i, j = symbols('i j', cls=Idx)
242 >>> Indexed('A', i, j).indices
243 (i, j)
245 """
246 return self.args[1:]
248 @property
249 def rank(self):
250 """
251 Returns the rank of the ``Indexed`` object.
253 Examples
254 ========
256 >>> from sympy import Indexed, Idx, symbols
257 >>> i, j, k, l, m = symbols('i:m', cls=Idx)
258 >>> Indexed('A', i, j).rank
259 2
260 >>> q = Indexed('A', i, j, k, l, m)
261 >>> q.rank
262 5
263 >>> q.rank == len(q.indices)
264 True
266 """
267 return len(self.args) - 1
269 @property
270 def shape(self):
271 """Returns a list with dimensions of each index.
273 Dimensions is a property of the array, not of the indices. Still, if
274 the ``IndexedBase`` does not define a shape attribute, it is assumed
275 that the ranges of the indices correspond to the shape of the array.
277 >>> from sympy import IndexedBase, Idx, symbols
278 >>> n, m = symbols('n m', integer=True)
279 >>> i = Idx('i', m)
280 >>> j = Idx('j', m)
281 >>> A = IndexedBase('A', shape=(n, n))
282 >>> B = IndexedBase('B')
283 >>> A[i, j].shape
284 (n, n)
285 >>> B[i, j].shape
286 (m, m)
287 """
289 if self.base.shape:
290 return self.base.shape
291 sizes = []
292 for i in self.indices:
293 upper = getattr(i, 'upper', None)
294 lower = getattr(i, 'lower', None)
295 if None in (upper, lower):
296 raise IndexException(filldedent("""
297 Range is not defined for all indices in: %s""" % self))
298 try:
299 size = upper - lower + 1
300 except TypeError:
301 raise IndexException(filldedent("""
302 Shape cannot be inferred from Idx with
303 undefined range: %s""" % self))
304 sizes.append(size)
305 return Tuple(*sizes)
307 @property
308 def ranges(self):
309 """Returns a list of tuples with lower and upper range of each index.
311 If an index does not define the data members upper and lower, the
312 corresponding slot in the list contains ``None`` instead of a tuple.
314 Examples
315 ========
317 >>> from sympy import Indexed,Idx, symbols
318 >>> Indexed('A', Idx('i', 2), Idx('j', 4), Idx('k', 8)).ranges
319 [(0, 1), (0, 3), (0, 7)]
320 >>> Indexed('A', Idx('i', 3), Idx('j', 3), Idx('k', 3)).ranges
321 [(0, 2), (0, 2), (0, 2)]
322 >>> x, y, z = symbols('x y z', integer=True)
323 >>> Indexed('A', x, y, z).ranges
324 [None, None, None]
326 """
327 ranges = []
328 sentinel = object()
329 for i in self.indices:
330 upper = getattr(i, 'upper', sentinel)
331 lower = getattr(i, 'lower', sentinel)
332 if sentinel not in (upper, lower):
333 ranges.append((lower, upper))
334 else:
335 ranges.append(None)
336 return ranges
338 def _sympystr(self, p):
339 indices = list(map(p.doprint, self.indices))
340 return "%s[%s]" % (p.doprint(self.base), ", ".join(indices))
342 @property
343 def free_symbols(self):
344 base_free_symbols = self.base.free_symbols
345 indices_free_symbols = {
346 fs for i in self.indices for fs in i.free_symbols}
347 if base_free_symbols:
348 return {self} | base_free_symbols | indices_free_symbols
349 else:
350 return indices_free_symbols
352 @property
353 def expr_free_symbols(self):
354 from sympy.utilities.exceptions import sympy_deprecation_warning
355 sympy_deprecation_warning("""
356 The expr_free_symbols property is deprecated. Use free_symbols to get
357 the free symbols of an expression.
358 """,
359 deprecated_since_version="1.9",
360 active_deprecations_target="deprecated-expr-free-symbols")
362 return {self}
365class IndexedBase(Expr, NotIterable):
366 """Represent the base or stem of an indexed object
368 The IndexedBase class represent an array that contains elements. The main purpose
369 of this class is to allow the convenient creation of objects of the Indexed
370 class. The __getitem__ method of IndexedBase returns an instance of
371 Indexed. Alone, without indices, the IndexedBase class can be used as a
372 notation for e.g. matrix equations, resembling what you could do with the
373 Symbol class. But, the IndexedBase class adds functionality that is not
374 available for Symbol instances:
376 - An IndexedBase object can optionally store shape information. This can
377 be used in to check array conformance and conditions for numpy
378 broadcasting. (TODO)
379 - An IndexedBase object implements syntactic sugar that allows easy symbolic
380 representation of array operations, using implicit summation of
381 repeated indices.
382 - The IndexedBase object symbolizes a mathematical structure equivalent
383 to arrays, and is recognized as such for code generation and automatic
384 compilation and wrapping.
386 >>> from sympy.tensor import IndexedBase, Idx
387 >>> from sympy import symbols
388 >>> A = IndexedBase('A'); A
389 A
390 >>> type(A)
391 <class 'sympy.tensor.indexed.IndexedBase'>
393 When an IndexedBase object receives indices, it returns an array with named
394 axes, represented by an Indexed object:
396 >>> i, j = symbols('i j', integer=True)
397 >>> A[i, j, 2]
398 A[i, j, 2]
399 >>> type(A[i, j, 2])
400 <class 'sympy.tensor.indexed.Indexed'>
402 The IndexedBase constructor takes an optional shape argument. If given,
403 it overrides any shape information in the indices. (But not the index
404 ranges!)
406 >>> m, n, o, p = symbols('m n o p', integer=True)
407 >>> i = Idx('i', m)
408 >>> j = Idx('j', n)
409 >>> A[i, j].shape
410 (m, n)
411 >>> B = IndexedBase('B', shape=(o, p))
412 >>> B[i, j].shape
413 (o, p)
415 Assumptions can be specified with keyword arguments the same way as for Symbol:
417 >>> A_real = IndexedBase('A', real=True)
418 >>> A_real.is_real
419 True
420 >>> A != A_real
421 True
423 Assumptions can also be inherited if a Symbol is used to initialize the IndexedBase:
425 >>> I = symbols('I', integer=True)
426 >>> C_inherit = IndexedBase(I)
427 >>> C_explicit = IndexedBase('I', integer=True)
428 >>> C_inherit == C_explicit
429 True
430 """
431 is_commutative = True
432 is_symbol = True
433 is_Atom = True
435 @staticmethod
436 def _set_assumptions(obj, assumptions):
437 """Set assumptions on obj, making sure to apply consistent values."""
438 tmp_asm_copy = assumptions.copy()
439 is_commutative = fuzzy_bool(assumptions.get('commutative', True))
440 assumptions['commutative'] = is_commutative
441 obj._assumptions = StdFactKB(assumptions)
442 obj._assumptions._generator = tmp_asm_copy # Issue #8873
444 def __new__(cls, label, shape=None, *, offset=S.Zero, strides=None, **kw_args):
445 from sympy.matrices.matrices import MatrixBase
446 from sympy.tensor.array.ndim_array import NDimArray
448 assumptions, kw_args = _filter_assumptions(kw_args)
449 if isinstance(label, str):
450 label = Symbol(label, **assumptions)
451 elif isinstance(label, Symbol):
452 assumptions = label._merge(assumptions)
453 elif isinstance(label, (MatrixBase, NDimArray)):
454 return label
455 elif isinstance(label, Iterable):
456 return _sympify(label)
457 else:
458 label = _sympify(label)
460 if is_sequence(shape):
461 shape = Tuple(*shape)
462 elif shape is not None:
463 shape = Tuple(shape)
465 if shape is not None:
466 obj = Expr.__new__(cls, label, shape)
467 else:
468 obj = Expr.__new__(cls, label)
469 obj._shape = shape
470 obj._offset = offset
471 obj._strides = strides
472 obj._name = str(label)
474 IndexedBase._set_assumptions(obj, assumptions)
475 return obj
477 @property
478 def name(self):
479 return self._name
481 def _hashable_content(self):
482 return super()._hashable_content() + tuple(sorted(self.assumptions0.items()))
484 @property
485 def assumptions0(self):
486 return {k: v for k, v in self._assumptions.items() if v is not None}
488 def __getitem__(self, indices, **kw_args):
489 if is_sequence(indices):
490 # Special case needed because M[*my_tuple] is a syntax error.
491 if self.shape and len(self.shape) != len(indices):
492 raise IndexException("Rank mismatch.")
493 return Indexed(self, *indices, **kw_args)
494 else:
495 if self.shape and len(self.shape) != 1:
496 raise IndexException("Rank mismatch.")
497 return Indexed(self, indices, **kw_args)
499 @property
500 def shape(self):
501 """Returns the shape of the ``IndexedBase`` object.
503 Examples
504 ========
506 >>> from sympy import IndexedBase, Idx
507 >>> from sympy.abc import x, y
508 >>> IndexedBase('A', shape=(x, y)).shape
509 (x, y)
511 Note: If the shape of the ``IndexedBase`` is specified, it will override
512 any shape information given by the indices.
514 >>> A = IndexedBase('A', shape=(x, y))
515 >>> B = IndexedBase('B')
516 >>> i = Idx('i', 2)
517 >>> j = Idx('j', 1)
518 >>> A[i, j].shape
519 (x, y)
520 >>> B[i, j].shape
521 (2, 1)
523 """
524 return self._shape
526 @property
527 def strides(self):
528 """Returns the strided scheme for the ``IndexedBase`` object.
530 Normally this is a tuple denoting the number of
531 steps to take in the respective dimension when traversing
532 an array. For code generation purposes strides='C' and
533 strides='F' can also be used.
535 strides='C' would mean that code printer would unroll
536 in row-major order and 'F' means unroll in column major
537 order.
539 """
541 return self._strides
543 @property
544 def offset(self):
545 """Returns the offset for the ``IndexedBase`` object.
547 This is the value added to the resulting index when the
548 2D Indexed object is unrolled to a 1D form. Used in code
549 generation.
551 Examples
552 ==========
553 >>> from sympy.printing import ccode
554 >>> from sympy.tensor import IndexedBase, Idx
555 >>> from sympy import symbols
556 >>> l, m, n, o = symbols('l m n o', integer=True)
557 >>> A = IndexedBase('A', strides=(l, m, n), offset=o)
558 >>> i, j, k = map(Idx, 'ijk')
559 >>> ccode(A[i, j, k])
560 'A[l*i + m*j + n*k + o]'
562 """
563 return self._offset
565 @property
566 def label(self):
567 """Returns the label of the ``IndexedBase`` object.
569 Examples
570 ========
572 >>> from sympy import IndexedBase
573 >>> from sympy.abc import x, y
574 >>> IndexedBase('A', shape=(x, y)).label
575 A
577 """
578 return self.args[0]
580 def _sympystr(self, p):
581 return p.doprint(self.label)
584class Idx(Expr):
585 """Represents an integer index as an ``Integer`` or integer expression.
587 There are a number of ways to create an ``Idx`` object. The constructor
588 takes two arguments:
590 ``label``
591 An integer or a symbol that labels the index.
592 ``range``
593 Optionally you can specify a range as either
595 * ``Symbol`` or integer: This is interpreted as a dimension. Lower and
596 upper bounds are set to ``0`` and ``range - 1``, respectively.
597 * ``tuple``: The two elements are interpreted as the lower and upper
598 bounds of the range, respectively.
600 Note: bounds of the range are assumed to be either integer or infinite (oo
601 and -oo are allowed to specify an unbounded range). If ``n`` is given as a
602 bound, then ``n.is_integer`` must not return false.
604 For convenience, if the label is given as a string it is automatically
605 converted to an integer symbol. (Note: this conversion is not done for
606 range or dimension arguments.)
608 Examples
609 ========
611 >>> from sympy import Idx, symbols, oo
612 >>> n, i, L, U = symbols('n i L U', integer=True)
614 If a string is given for the label an integer ``Symbol`` is created and the
615 bounds are both ``None``:
617 >>> idx = Idx('qwerty'); idx
618 qwerty
619 >>> idx.lower, idx.upper
620 (None, None)
622 Both upper and lower bounds can be specified:
624 >>> idx = Idx(i, (L, U)); idx
625 i
626 >>> idx.lower, idx.upper
627 (L, U)
629 When only a single bound is given it is interpreted as the dimension
630 and the lower bound defaults to 0:
632 >>> idx = Idx(i, n); idx.lower, idx.upper
633 (0, n - 1)
634 >>> idx = Idx(i, 4); idx.lower, idx.upper
635 (0, 3)
636 >>> idx = Idx(i, oo); idx.lower, idx.upper
637 (0, oo)
639 """
641 is_integer = True
642 is_finite = True
643 is_real = True
644 is_symbol = True
645 is_Atom = True
646 _diff_wrt = True
648 def __new__(cls, label, range=None, **kw_args):
650 if isinstance(label, str):
651 label = Symbol(label, integer=True)
652 label, range = list(map(sympify, (label, range)))
654 if label.is_Number:
655 if not label.is_integer:
656 raise TypeError("Index is not an integer number.")
657 return label
659 if not label.is_integer:
660 raise TypeError("Idx object requires an integer label.")
662 elif is_sequence(range):
663 if len(range) != 2:
664 raise ValueError(filldedent("""
665 Idx range tuple must have length 2, but got %s""" % len(range)))
666 for bound in range:
667 if (bound.is_integer is False and bound is not S.Infinity
668 and bound is not S.NegativeInfinity):
669 raise TypeError("Idx object requires integer bounds.")
670 args = label, Tuple(*range)
671 elif isinstance(range, Expr):
672 if range is not S.Infinity and fuzzy_not(range.is_integer):
673 raise TypeError("Idx object requires an integer dimension.")
674 args = label, Tuple(0, range - 1)
675 elif range:
676 raise TypeError(filldedent("""
677 The range must be an ordered iterable or
678 integer SymPy expression."""))
679 else:
680 args = label,
682 obj = Expr.__new__(cls, *args, **kw_args)
683 obj._assumptions["finite"] = True
684 obj._assumptions["real"] = True
685 return obj
687 @property
688 def label(self):
689 """Returns the label (Integer or integer expression) of the Idx object.
691 Examples
692 ========
694 >>> from sympy import Idx, Symbol
695 >>> x = Symbol('x', integer=True)
696 >>> Idx(x).label
697 x
698 >>> j = Symbol('j', integer=True)
699 >>> Idx(j).label
700 j
701 >>> Idx(j + 1).label
702 j + 1
704 """
705 return self.args[0]
707 @property
708 def lower(self):
709 """Returns the lower bound of the ``Idx``.
711 Examples
712 ========
714 >>> from sympy import Idx
715 >>> Idx('j', 2).lower
716 0
717 >>> Idx('j', 5).lower
718 0
719 >>> Idx('j').lower is None
720 True
722 """
723 try:
724 return self.args[1][0]
725 except IndexError:
726 return
728 @property
729 def upper(self):
730 """Returns the upper bound of the ``Idx``.
732 Examples
733 ========
735 >>> from sympy import Idx
736 >>> Idx('j', 2).upper
737 1
738 >>> Idx('j', 5).upper
739 4
740 >>> Idx('j').upper is None
741 True
743 """
744 try:
745 return self.args[1][1]
746 except IndexError:
747 return
749 def _sympystr(self, p):
750 return p.doprint(self.label)
752 @property
753 def name(self):
754 return self.label.name if self.label.is_Symbol else str(self.label)
756 @property
757 def free_symbols(self):
758 return {self}
761@dispatch(Idx, Idx)
762def _eval_is_ge(lhs, rhs): # noqa:F811
764 other_upper = rhs if rhs.upper is None else rhs.upper
765 other_lower = rhs if rhs.lower is None else rhs.lower
767 if lhs.lower is not None and (lhs.lower >= other_upper) == True:
768 return True
769 if lhs.upper is not None and (lhs.upper < other_lower) == True:
770 return False
771 return None
774@dispatch(Idx, Number) # type:ignore
775def _eval_is_ge(lhs, rhs): # noqa:F811
777 other_upper = rhs
778 other_lower = rhs
780 if lhs.lower is not None and (lhs.lower >= other_upper) == True:
781 return True
782 if lhs.upper is not None and (lhs.upper < other_lower) == True:
783 return False
784 return None
787@dispatch(Number, Idx) # type:ignore
788def _eval_is_ge(lhs, rhs): # noqa:F811
790 other_upper = lhs
791 other_lower = lhs
793 if rhs.upper is not None and (rhs.upper <= other_lower) == True:
794 return True
795 if rhs.lower is not None and (rhs.lower > other_upper) == True:
796 return False
797 return None