CODEWITHSUNDEEP
pythonsympy
MB-F
MB-F

Reputation: 23637

List/Array of matrices in sympy

I want to create the following expression in sympy:

\sum_{k=0}^p \mathbf{M}\mathbf{A}^{(k)}\mathbf{U}

M, U and A^(k) are square matrices of size m.

So far, I managed to do this for a given p:

import sympy

m = sympy.symbols('m', integer=True)
p = 3

A = [sympy.MatrixSymbol('A^({})'.format(k), m, m) for k in range(p)]
M = sympy.MatrixSymbol('M', m, m)
U = sympy.MatrixSymbol('U', m, m)

expr = M*A[0]*U
for k in range(1, p):
    expr += M*A[k]*U
expr

enter image description here

However, I do not want to set p to a fixed number. Instead I want p to be a symbol, just like m. How can this be achieved in sympy?

Upvotes: 1

Views: 585

Answers (1)

MB-F
MB-F

Reputation: 23637

There may be a better way. For now, I managed to achieve what I want by defining A as a function rather than a list:

m, p, k = symbols('m, p, k')

M = sympy.MatrixSymbol('M', m, m)
U = sympy.MatrixSymbol('U', m, m)

class A(sympy.Function):
    @classmethod
    def eval(cls, k):
        return sympy.MatrixSymbol('A^({})'.format(k), m, m)

sympy.Sum(M * A(k) * U, (k, 1, p))

Upvotes: 2

Related Questions