How to expand matrix expressions in SymPy?

Question

In SymPy, I am trying to perform a matrix multiplication and expand it afterwards. However, SymPy does not seem to support the expansion of matrix expressions. For example, here is the 4th order Runge-Kutta (RK4) for matrices:

from sympy import init_session
init_session()
from sympy import *

A = MatrixSymbol('A', 3, 3)
x = MatrixSymbol('x', 3, 1)
dt = symbols('dt')

k1 = A*x
k2 = A*(x + S(1)/2*k1*dt)
k3 = A*(x + S(1)/2*k2*dt)
k4 = A*(x + k3*dt)
final = dt*S(1)/6*(k1 + 2*k2 + 2*k3 + k4)
final.expand()

which produces the result

Traceback (most recent call last)
<ipython-input-38-b3ff67883c61> in <module>()
     12 final = dt*1/6*(k1+2*k2+2*k3+k4)
     13 
---> 14 final.expand()

AttributeError: 'MatMul' object has no attribute 'expand'

I hope the expression can be expanded just like the scalar variant:

A,x,dt = symbols('A x dt')
k1 = A*x
k2 = A*(x+k1*dt*S(1)/2)
k3 = A*(x+k2*dt*S(1)/2)
k4 = A*(x+k3*dt)
final = x+dt*(S(1)/6)*(k1+k2+k3+k4)
collect(expand((final)),x)

with the result:

x*(A**4*dt**4/24 + A**3*dt**3/8 + A**2*dt**2/3 + 2*A*dt/3 + 1)

Is it possible to alter a matrix expression likewise?

nicoguaro's answer takes the error away, but expands the the whole expression into one matrix. As illustrated with the scalar example, not what I'm looking for.

Answer 1

Matrix(final) creates and explicit Matrix with your individual equations. It may be convenient to leave them in the matrix so that whatever you do to one entry can be done to all. To do so, define the manipulation that you want to do as a function argument for applyfunc:

>>> ex = Matrix(final)
>>> ex = ex.applyfunc(expand)
>>> ex = ex.applyfunc(lambda i: collect(i, dt))
...

For printing economy I use Matrices with compact symbol entries to compute the same and then run cse over the simplified matrix to give:

>>> A = Matrix(3, 3, var('a:3:3'))
>>> x = Matrix(3, 1, var('x:3'))
>>> dt = symbols('dt')
>>> k1 = A*x
>>> k2 = A*(x + S(1)/2*k1*dt)
>>> k3 = A*(x + S(1)/2*k2*dt)
>>> k4 = A*(x + k3*dt)
>>> final = dt*S(1)/6*(k1 + 2*k2 + 2*k3 + k4)
>>> eqi = Matrix(final)
>>> print(cse(eqi.applyfunc(simplify)))
([(x3, 4*x0), (x4, a01*x1), (x5, a02*x2), (x6, dt*(a00*x0 + x4 + x5) +
2*x0), (x7, a00*x6), (x8, a11*x1), (x9, a12*x2), (x10, dt*(a10*x0 + x8
+ x9) + 2*x1), (x11, a01*x10), (x12, a21*x1), (x13, a22*x2), (x14,
dt*(a20*x0 + x12 + x13) + 2*x2), (x15, a02*x14), (x16, dt*(x11 + x15 +
x7) + x3), (x17, a00*x16), (x18, 4*x1), (x19, a10*x6), (x20, a11*x10),
(x21, a12*x14), (x22, dt*(x19 + x20 + x21) + x18), (x23, a01*x22),
(x24, 4*x2), (x25, a20*x6), (x26, a21*x10), (x27, a22*x14), (x28,
dt*(x25 + x26 + x27) + x24), (x29, a02*x28), (x30, dt*(x17 + x23 +
x29) + x3), (x31, a10*x16), (x32, a11*x22), (x33, a12*x28), (x34,
dt*(x31 + x32 + x33) + x18), (x35, a20*x16), (x36, a21*x22), (x37,
a22*x28), (x38, dt*(x35 + x36 + x37) + x24), (x39, dt/24)], [Matrix([
[   x39*(a00*x3 + a00*x30 + a01*x34 + a02*x38 + 4*x11 + 4*x15 + 2*x17
+ 2*x23 + 2*x29 + 4*x4 + 4*x5 + 4*x7)], [  x39*(a10*x3 + a10*x30 +
a11*x34 + a12*x38 + 4*x19 + 4*x20 + 4*x21 + 2*x31 + 2*x32 + 2*x33 +
4*x8 + 4*x9)], [x39*(a20*x3 + a20*x30 + a21*x34 + a22*x38 + 4*x12 +
4*x13 + 4*x25 + 4*x26 + 4*x27 + 2*x35 + 2*x36 + 2*x37)]])])

Limited support of noncommutative expressions is also available and can help in this situation:

>>> A, x = symbols("A x", commutative=False)
>>> dt = symbols('dt')
>>> k1 = A*x
>>> k2 = A*(x + S(1)/2*k1*dt)
>>> k3 = A*(x + S(1)/2*k2*dt)
>>> k4 = A*(x + k3*dt)
>>> final = dt*S(1)/6*(k1 + 2*k2 + 2*k3 + k4)
>>> final.expand()
dt**4*A**4*x/24 + dt**3*A**3*x/6 + dt**2*A**2*x/2 + dt*A*x
>>> factor(_)
dt*A*(dt**3*A**3/24 + dt**2*A**2/6 + dt*A/2 + 1)*x

But not all simplification routines are nc aware (and this is a known issue):

>>> collect(final,x)
Traceback (most recent call last):
...
AttributeError: Can not collect noncommutative symbol