Coverage for /usr/lib/python3/dist-packages/sympy/tensor/index_methods.py: 11%
133 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""Module with functions operating on IndexedBase, Indexed and Idx objects
3 - Check shape conformance
4 - Determine indices in resulting expression
6 etc.
8 Methods in this module could be implemented by calling methods on Expr
9 objects instead. When things stabilize this could be a useful
10 refactoring.
11"""
13from functools import reduce
15from sympy.core.function import Function
16from sympy.functions import exp, Piecewise
17from sympy.tensor.indexed import Idx, Indexed
18from sympy.utilities import sift
20from collections import OrderedDict
22class IndexConformanceException(Exception):
23 pass
25def _unique_and_repeated(inds):
26 """
27 Returns the unique and repeated indices. Also note, from the examples given below
28 that the order of indices is maintained as given in the input.
30 Examples
31 ========
33 >>> from sympy.tensor.index_methods import _unique_and_repeated
34 >>> _unique_and_repeated([2, 3, 1, 3, 0, 4, 0])
35 ([2, 1, 4], [3, 0])
36 """
37 uniq = OrderedDict()
38 for i in inds:
39 if i in uniq:
40 uniq[i] = 0
41 else:
42 uniq[i] = 1
43 return sift(uniq, lambda x: uniq[x], binary=True)
45def _remove_repeated(inds):
46 """
47 Removes repeated objects from sequences
49 Returns a set of the unique objects and a tuple of all that have been
50 removed.
52 Examples
53 ========
55 >>> from sympy.tensor.index_methods import _remove_repeated
56 >>> l1 = [1, 2, 3, 2]
57 >>> _remove_repeated(l1)
58 ({1, 3}, (2,))
60 """
61 u, r = _unique_and_repeated(inds)
62 return set(u), tuple(r)
65def _get_indices_Mul(expr, return_dummies=False):
66 """Determine the outer indices of a Mul object.
68 Examples
69 ========
71 >>> from sympy.tensor.index_methods import _get_indices_Mul
72 >>> from sympy.tensor.indexed import IndexedBase, Idx
73 >>> i, j, k = map(Idx, ['i', 'j', 'k'])
74 >>> x = IndexedBase('x')
75 >>> y = IndexedBase('y')
76 >>> _get_indices_Mul(x[i, k]*y[j, k])
77 ({i, j}, {})
78 >>> _get_indices_Mul(x[i, k]*y[j, k], return_dummies=True)
79 ({i, j}, {}, (k,))
81 """
83 inds = list(map(get_indices, expr.args))
84 inds, syms = list(zip(*inds))
86 inds = list(map(list, inds))
87 inds = list(reduce(lambda x, y: x + y, inds))
88 inds, dummies = _remove_repeated(inds)
90 symmetry = {}
91 for s in syms:
92 for pair in s:
93 if pair in symmetry:
94 symmetry[pair] *= s[pair]
95 else:
96 symmetry[pair] = s[pair]
98 if return_dummies:
99 return inds, symmetry, dummies
100 else:
101 return inds, symmetry
104def _get_indices_Pow(expr):
105 """Determine outer indices of a power or an exponential.
107 A power is considered a universal function, so that the indices of a Pow is
108 just the collection of indices present in the expression. This may be
109 viewed as a bit inconsistent in the special case:
111 x[i]**2 = x[i]*x[i] (1)
113 The above expression could have been interpreted as the contraction of x[i]
114 with itself, but we choose instead to interpret it as a function
116 lambda y: y**2
118 applied to each element of x (a universal function in numpy terms). In
119 order to allow an interpretation of (1) as a contraction, we need
120 contravariant and covariant Idx subclasses. (FIXME: this is not yet
121 implemented)
123 Expressions in the base or exponent are subject to contraction as usual,
124 but an index that is present in the exponent, will not be considered
125 contractable with its own base. Note however, that indices in the same
126 exponent can be contracted with each other.
128 Examples
129 ========
131 >>> from sympy.tensor.index_methods import _get_indices_Pow
132 >>> from sympy import Pow, exp, IndexedBase, Idx
133 >>> A = IndexedBase('A')
134 >>> x = IndexedBase('x')
135 >>> i, j, k = map(Idx, ['i', 'j', 'k'])
136 >>> _get_indices_Pow(exp(A[i, j]*x[j]))
137 ({i}, {})
138 >>> _get_indices_Pow(Pow(x[i], x[i]))
139 ({i}, {})
140 >>> _get_indices_Pow(Pow(A[i, j]*x[j], x[i]))
141 ({i}, {})
143 """
144 base, exp = expr.as_base_exp()
145 binds, bsyms = get_indices(base)
146 einds, esyms = get_indices(exp)
148 inds = binds | einds
150 # FIXME: symmetries from power needs to check special cases, else nothing
151 symmetries = {}
153 return inds, symmetries
156def _get_indices_Add(expr):
157 """Determine outer indices of an Add object.
159 In a sum, each term must have the same set of outer indices. A valid
160 expression could be
162 x(i)*y(j) - x(j)*y(i)
164 But we do not allow expressions like:
166 x(i)*y(j) - z(j)*z(j)
168 FIXME: Add support for Numpy broadcasting
170 Examples
171 ========
173 >>> from sympy.tensor.index_methods import _get_indices_Add
174 >>> from sympy.tensor.indexed import IndexedBase, Idx
175 >>> i, j, k = map(Idx, ['i', 'j', 'k'])
176 >>> x = IndexedBase('x')
177 >>> y = IndexedBase('y')
178 >>> _get_indices_Add(x[i] + x[k]*y[i, k])
179 ({i}, {})
181 """
183 inds = list(map(get_indices, expr.args))
184 inds, syms = list(zip(*inds))
186 # allow broadcast of scalars
187 non_scalars = [x for x in inds if x != set()]
188 if not non_scalars:
189 return set(), {}
191 if not all(x == non_scalars[0] for x in non_scalars[1:]):
192 raise IndexConformanceException("Indices are not consistent: %s" % expr)
193 if not reduce(lambda x, y: x != y or y, syms):
194 symmetries = syms[0]
195 else:
196 # FIXME: search for symmetries
197 symmetries = {}
199 return non_scalars[0], symmetries
202def get_indices(expr):
203 """Determine the outer indices of expression ``expr``
205 By *outer* we mean indices that are not summation indices. Returns a set
206 and a dict. The set contains outer indices and the dict contains
207 information about index symmetries.
209 Examples
210 ========
212 >>> from sympy.tensor.index_methods import get_indices
213 >>> from sympy import symbols
214 >>> from sympy.tensor import IndexedBase
215 >>> x, y, A = map(IndexedBase, ['x', 'y', 'A'])
216 >>> i, j, a, z = symbols('i j a z', integer=True)
218 The indices of the total expression is determined, Repeated indices imply a
219 summation, for instance the trace of a matrix A:
221 >>> get_indices(A[i, i])
222 (set(), {})
224 In the case of many terms, the terms are required to have identical
225 outer indices. Else an IndexConformanceException is raised.
227 >>> get_indices(x[i] + A[i, j]*y[j])
228 ({i}, {})
230 :Exceptions:
232 An IndexConformanceException means that the terms ar not compatible, e.g.
234 >>> get_indices(x[i] + y[j]) #doctest: +SKIP
235 (...)
236 IndexConformanceException: Indices are not consistent: x(i) + y(j)
238 .. warning::
239 The concept of *outer* indices applies recursively, starting on the deepest
240 level. This implies that dummies inside parenthesis are assumed to be
241 summed first, so that the following expression is handled gracefully:
243 >>> get_indices((x[i] + A[i, j]*y[j])*x[j])
244 ({i, j}, {})
246 This is correct and may appear convenient, but you need to be careful
247 with this as SymPy will happily .expand() the product, if requested. The
248 resulting expression would mix the outer ``j`` with the dummies inside
249 the parenthesis, which makes it a different expression. To be on the
250 safe side, it is best to avoid such ambiguities by using unique indices
251 for all contractions that should be held separate.
253 """
254 # We call ourself recursively to determine indices of sub expressions.
256 # break recursion
257 if isinstance(expr, Indexed):
258 c = expr.indices
259 inds, dummies = _remove_repeated(c)
260 return inds, {}
261 elif expr is None:
262 return set(), {}
263 elif isinstance(expr, Idx):
264 return {expr}, {}
265 elif expr.is_Atom:
266 return set(), {}
269 # recurse via specialized functions
270 else:
271 if expr.is_Mul:
272 return _get_indices_Mul(expr)
273 elif expr.is_Add:
274 return _get_indices_Add(expr)
275 elif expr.is_Pow or isinstance(expr, exp):
276 return _get_indices_Pow(expr)
278 elif isinstance(expr, Piecewise):
279 # FIXME: No support for Piecewise yet
280 return set(), {}
281 elif isinstance(expr, Function):
282 # Support ufunc like behaviour by returning indices from arguments.
283 # Functions do not interpret repeated indices across arguments
284 # as summation
285 ind0 = set()
286 for arg in expr.args:
287 ind, sym = get_indices(arg)
288 ind0 |= ind
289 return ind0, sym
291 # this test is expensive, so it should be at the end
292 elif not expr.has(Indexed):
293 return set(), {}
294 raise NotImplementedError(
295 "FIXME: No specialized handling of type %s" % type(expr))
298def get_contraction_structure(expr):
299 """Determine dummy indices of ``expr`` and describe its structure
301 By *dummy* we mean indices that are summation indices.
303 The structure of the expression is determined and described as follows:
305 1) A conforming summation of Indexed objects is described with a dict where
306 the keys are summation indices and the corresponding values are sets
307 containing all terms for which the summation applies. All Add objects
308 in the SymPy expression tree are described like this.
310 2) For all nodes in the SymPy expression tree that are *not* of type Add, the
311 following applies:
313 If a node discovers contractions in one of its arguments, the node
314 itself will be stored as a key in the dict. For that key, the
315 corresponding value is a list of dicts, each of which is the result of a
316 recursive call to get_contraction_structure(). The list contains only
317 dicts for the non-trivial deeper contractions, omitting dicts with None
318 as the one and only key.
320 .. Note:: The presence of expressions among the dictionary keys indicates
321 multiple levels of index contractions. A nested dict displays nested
322 contractions and may itself contain dicts from a deeper level. In
323 practical calculations the summation in the deepest nested level must be
324 calculated first so that the outer expression can access the resulting
325 indexed object.
327 Examples
328 ========
330 >>> from sympy.tensor.index_methods import get_contraction_structure
331 >>> from sympy import default_sort_key
332 >>> from sympy.tensor import IndexedBase, Idx
333 >>> x, y, A = map(IndexedBase, ['x', 'y', 'A'])
334 >>> i, j, k, l = map(Idx, ['i', 'j', 'k', 'l'])
335 >>> get_contraction_structure(x[i]*y[i] + A[j, j])
336 {(i,): {x[i]*y[i]}, (j,): {A[j, j]}}
337 >>> get_contraction_structure(x[i]*y[j])
338 {None: {x[i]*y[j]}}
340 A multiplication of contracted factors results in nested dicts representing
341 the internal contractions.
343 >>> d = get_contraction_structure(x[i, i]*y[j, j])
344 >>> sorted(d.keys(), key=default_sort_key)
345 [None, x[i, i]*y[j, j]]
347 In this case, the product has no contractions:
349 >>> d[None]
350 {x[i, i]*y[j, j]}
352 Factors are contracted "first":
354 >>> sorted(d[x[i, i]*y[j, j]], key=default_sort_key)
355 [{(i,): {x[i, i]}}, {(j,): {y[j, j]}}]
357 A parenthesized Add object is also returned as a nested dictionary. The
358 term containing the parenthesis is a Mul with a contraction among the
359 arguments, so it will be found as a key in the result. It stores the
360 dictionary resulting from a recursive call on the Add expression.
362 >>> d = get_contraction_structure(x[i]*(y[i] + A[i, j]*x[j]))
363 >>> sorted(d.keys(), key=default_sort_key)
364 [(A[i, j]*x[j] + y[i])*x[i], (i,)]
365 >>> d[(i,)]
366 {(A[i, j]*x[j] + y[i])*x[i]}
367 >>> d[x[i]*(A[i, j]*x[j] + y[i])]
368 [{None: {y[i]}, (j,): {A[i, j]*x[j]}}]
370 Powers with contractions in either base or exponent will also be found as
371 keys in the dictionary, mapping to a list of results from recursive calls:
373 >>> d = get_contraction_structure(A[j, j]**A[i, i])
374 >>> d[None]
375 {A[j, j]**A[i, i]}
376 >>> nested_contractions = d[A[j, j]**A[i, i]]
377 >>> nested_contractions[0]
378 {(j,): {A[j, j]}}
379 >>> nested_contractions[1]
380 {(i,): {A[i, i]}}
382 The description of the contraction structure may appear complicated when
383 represented with a string in the above examples, but it is easy to iterate
384 over:
386 >>> from sympy import Expr
387 >>> for key in d:
388 ... if isinstance(key, Expr):
389 ... continue
390 ... for term in d[key]:
391 ... if term in d:
392 ... # treat deepest contraction first
393 ... pass
394 ... # treat outermost contactions here
396 """
398 # We call ourself recursively to inspect sub expressions.
400 if isinstance(expr, Indexed):
401 junk, key = _remove_repeated(expr.indices)
402 return {key or None: {expr}}
403 elif expr.is_Atom:
404 return {None: {expr}}
405 elif expr.is_Mul:
406 junk, junk, key = _get_indices_Mul(expr, return_dummies=True)
407 result = {key or None: {expr}}
408 # recurse on every factor
409 nested = []
410 for fac in expr.args:
411 facd = get_contraction_structure(fac)
412 if not (None in facd and len(facd) == 1):
413 nested.append(facd)
414 if nested:
415 result[expr] = nested
416 return result
417 elif expr.is_Pow or isinstance(expr, exp):
418 # recurse in base and exp separately. If either has internal
419 # contractions we must include ourselves as a key in the returned dict
420 b, e = expr.as_base_exp()
421 dbase = get_contraction_structure(b)
422 dexp = get_contraction_structure(e)
424 dicts = []
425 for d in dbase, dexp:
426 if not (None in d and len(d) == 1):
427 dicts.append(d)
428 result = {None: {expr}}
429 if dicts:
430 result[expr] = dicts
431 return result
432 elif expr.is_Add:
433 # Note: we just collect all terms with identical summation indices, We
434 # do nothing to identify equivalent terms here, as this would require
435 # substitutions or pattern matching in expressions of unknown
436 # complexity.
437 result = {}
438 for term in expr.args:
439 # recurse on every term
440 d = get_contraction_structure(term)
441 for key in d:
442 if key in result:
443 result[key] |= d[key]
444 else:
445 result[key] = d[key]
446 return result
448 elif isinstance(expr, Piecewise):
449 # FIXME: No support for Piecewise yet
450 return {None: expr}
451 elif isinstance(expr, Function):
452 # Collect non-trivial contraction structures in each argument
453 # We do not report repeated indices in separate arguments as a
454 # contraction
455 deeplist = []
456 for arg in expr.args:
457 deep = get_contraction_structure(arg)
458 if not (None in deep and len(deep) == 1):
459 deeplist.append(deep)
460 d = {None: {expr}}
461 if deeplist:
462 d[expr] = deeplist
463 return d
465 # this test is expensive, so it should be at the end
466 elif not expr.has(Indexed):
467 return {None: {expr}}
468 raise NotImplementedError(
469 "FIXME: No specialized handling of type %s" % type(expr))