CODEWITHSUNDEEP
sympy
Mark
Mark

Reputation: 5499

Getting used indices of indexed terms in a sympy sum

Say I have variables x and y which are indexed

from sympy.tensor import IndexedBase
x = IndexedBase('x')`

And I have an expression like e = x[1] y[2] + x[5] y[10] . I want to find all indices used by each of x and y. I'm looking for a function which might work like this: e.indices(y) = [2, 10] and e.indicies(x) = [1, 5]

Is there a way I can iterate through the terms x[i] y[j]? And if so, is there a way to split a product into terms and for each of those pull out which letter is being used and which index appears?

Upvotes: 0

Views: 263

Answers (1)

smichr
smichr

Reputation: 19057

The following should get you headed in the right direction:

>>> from sympy.tensor import IndexedBase, Indexed
>>> from sympy import sift
>>> x = IndexedBase('x')
>>> y = IndexedBase('y')
>>> e = x[1]* y[2] + x[5]* y[10]
>>> e.atoms(IndexedBase)
set([y, x])
>>> e.atoms(Indexed)
set([x[5], y[10], x[1], y[2]])
>>> sifted = sift(_,lambda i: i.base)
>>> sifted[x]
[x[5], x[1]]
>>> sifted[y]
[y[10], y[2]]
>>> [i.indices for i in _]
[(10,), (2,)]
>>> flatten(_)
[10, 2]

Upvotes: 1

Related Questions