I have to compute the following sum, S defined as:
I have tried the following as a function:
import numpy as np
a = np.array([0.4,2.6,3, 1.2, 3.4])
b = np.array([-5,7,2,1.1,1.8])
c = np.array([3.3,30,15,0.4,28])
t = np.linspace(0, np.pi, 600)
def S(t):
return np.sum([a[i]*np.cos(b[i]*t + c[i]) for i in range(5)], axis=0)
This works just fine. But I wonder if there is pure numpy version that uses broadcasting without relying on Python's list comprehension?
I have tried:
def S(t):
return np.sum(a*np.cos(b*t + c), axis=0)
When I compue S(t) I get the following error:
...
ValueError: operands could not be broadcast together with shapes (5,) (600,)
How can I get this to work properly?
How to do it?
For broadcasting you need ones in dimensions you want to have "copied" values. So it works like this:
import numpy as np
a = np.array([0.4,2.6,3, 1.2, 3.4])
b = np.array([-5,7,2,1.1,1.8])
c = np.array([3.3,30,15,0.4,28])
t = np.linspace(0, np.pi, 600).reshape(-1,1)
np.sum(a*np.cos(b*t+c), axis=1)
How to check if you did the right thing?
One neat hack you can do is to put sympy.symbols into your arrays and do the same thing. Then you obtain formulas for what it calculated for you:
import sympy
a = np.array(sympy.symbols([f"a{i}" for i in range(5)]))
b = np.array(sympy.symbols([f"b{i}" for i in range(5)]))
c = np.array(sympy.symbols([f"c{i}" for i in range(5)]))
t = np.array(sympy.symbols([f"t{i}" for i in range(10)])).reshape(-1,1)
cos = np.vectorize(sympy.cos)
np.sum(a*cos(b*t+c), axis=1)
this gives you
array([a0*cos(b0*t0 + c0) + a1*cos(b1*t0 + c1) + a2*cos(b2*t0 + c2) + a3*cos(b3*t0 + c3) + a4*cos(b4*t0 + c4),
a0*cos(b0*t1 + c0) + a1*cos(b1*t1 + c1) + a2*cos(b2*t1 + c2) + a3*cos(b3*t1 + c3) + a4*cos(b4*t1 + c4),
a0*cos(b0*t2 + c0) + a1*cos(b1*t2 + c1) + a2*cos(b2*t2 + c2) + a3*cos(b3*t2 + c3) + a4*cos(b4*t2 + c4),
a0*cos(b0*t3 + c0) + a1*cos(b1*t3 + c1) + a2*cos(b2*t3 + c2) + a3*cos(b3*t3 + c3) + a4*cos(b4*t3 + c4),
a0*cos(b0*t4 + c0) + a1*cos(b1*t4 + c1) + a2*cos(b2*t4 + c2) + a3*cos(b3*t4 + c3) + a4*cos(b4*t4 + c4),
a0*cos(b0*t5 + c0) + a1*cos(b1*t5 + c1) + a2*cos(b2*t5 + c2) + a3*cos(b3*t5 + c3) + a4*cos(b4*t5 + c4),
a0*cos(b0*t6 + c0) + a1*cos(b1*t6 + c1) + a2*cos(b2*t6 + c2) + a3*cos(b3*t6 + c3) + a4*cos(b4*t6 + c4),
a0*cos(b0*t7 + c0) + a1*cos(b1*t7 + c1) + a2*cos(b2*t7 + c2) + a3*cos(b3*t7 + c3) + a4*cos(b4*t7 + c4),
a0*cos(b0*t8 + c0) + a1*cos(b1*t8 + c1) + a2*cos(b2*t8 + c2) + a3*cos(b3*t8 + c3) + a4*cos(b4*t8 + c4),
a0*cos(b0*t9 + c0) + a1*cos(b1*t9 + c1) + a2*cos(b2*t9 + c2) + a3*cos(b3*t9 + c3) + a4*cos(b4*t9 + c4)],
dtype=object)
Notice that this couldn't work with numpy's cosine since that only works for floats. But
cos = np.vectorize(sympy.cos)
gives you a version of sympys cosine that works on arrays elementwise.
I think I found the answer by adding a newaxis to a, b, and c arrays as follows:
def S(t):
return np.sum(a[:,None]*cos(b[:,None]*t + c[:,None]), axis=0)
or
def S(t):
return np.sum(a[:,np.newaxis]*cos(b[:,np.newaxis]*t + c[:,np.newaxis]), axis=0)
Edit:
As suggested by Michael in the comment, a better solution, would be to add a newaxis to t instead (note that in this case the sum is carried out over axis=1):
def S(t):
return np.sum(a*cos(b*t[:,None] + c), axis=1)
Related
I’m experiencing a weird behaviour with this code. I would expect that w_numba == w but this is not the case. if Nx <6 the results are the same but for higher number the numba compiled method returns nan.
Could you please help me?
Many thanks
Regards
import numpy as np
from numba import njit
#njit
def dp_formula_with_numba(a: np.ndarray, b: float, x: np.ndarray):
""" formula dp """
term1 = np.arcsinh((b - 1) / (a ** .5))
return ((1 - b) * np.arcsinh(b / (a ** .5)) - ((b ** 2 - 2 * b + a + 1) ** .5) + ((b ** 2 + a) ** .5) +
term1 * (b - 1)) * ((1 - x ** 2) ** .5)
def dp_formula(a, b, x):
""" formula dp """
term1 = np.arcsinh((b - 1) / (a ** .5))
return ((1 - b) * np.arcsinh(b / (a ** .5)) - ((b ** 2 - 2 * b + a + 1) ** .5) + ((b ** 2 + a) ** .5) +
term1 * (b - 1)) * ((1 - x ** 2) ** .5)
Nx = 8
Nint = int(Nx/2)
XX1 = np.arange(0, Nint + 1) / Nint
XX1[0] = 0.0025
XX1[Nint] = 0.9999
b_int = 0
a_int = (0.25 * XX1) ** 2
print("a: " + str(a_int) + " dtype:" + str(a_int.dtype) + " type:" + str(type(a_int)))
print("b: " + str(b_int) + " type:" + str(type(b_int)))
print("x: " + str(XX1) + " dtype:" + str(XX1.dtype) + " type:" + str(type(XX1)))
w_numba = dp_formula_with_numba(a_int, b_int, XX1)
w = dp_formula(a_int, b_int, XX1)
print("")
print('With numba: ' + str(w_numba))
print('With numpy: ' + str(w))
Results:
a_int: [3.90625000e-07 3.90625000e-03 1.56250000e-02 3.51562500e-02 6.24875006e-02] dtype:float64 type:<class 'numpy.ndarray'>
b_int: 0 type:<class 'int'>
XX1: [0.0025 0.25 0.5 0.75 0.9999] dtype:float64 type:<class 'numpy.ndarray'>
With numba: [7.07150889 2.4470088 nan nan 0.0185825 ]
With numpy: [7.07150889 2.4470088 1.6399837 1.0224987 0.0185825 ]
I did an even simpler test:
import numpy as np
from numba import njit
#njit
def testing_with_numba(a: np.ndarray, b: float):
""" formula dp """
q = a**.5
print((b - 1) / q)
print((b - 1) / (a**.5))
a_int = np.arange(4)+1
b_int = 0
print("a_int: " + str(a_int) + " dtype:" + str(a_int.dtype) + " type:" + str(type(a_int)))
print("b_int: " + str(b_int) + " type:" + str(type(b_int)))
testing_with_numba(a_int, b_int)
and this is what gets printed:
a_int: [1 2 3 4] dtype:int32 type:<class 'numpy.ndarray'>
b_int: 0 type:<class 'int'>
[-1. -0.70710678 -0.57735027 -0.5 ]
[-1. -0.70710678 nan nan]
I tried with both python3.6 + numba 0.53.1 + numpy 1.19.5, and python3.8 + numba 0.51.2 + numpy 1.19.2 on Windows 7 and got the same results.
I then ran the same code in the same environment (Anaconda3-2020.11-Windows-x86_64, python3.8 + numba 0.51.2 + numpy 1.19.2 ) but on windows 10 and it works fine....
how to solve a set of expressions in Python, not numerical calcualtion but purely analytic way?
I checked several functions of sympy which seems unable to do the work.
Assuming we are going to find the expression for x1,x2,x3 and x4 with following equations:
x1 = a*s+b*x2 + c
x2 = (m+n) * x3 + d
x3 = l*s + k*x4 + e
x4 = i*s + j*x1 + f
where x1,x2,x3,x4,a,b,c,d,e,f,i,j,k,l,m,n are all symbols.
How to reach the result for such problems in Python.
First get the "paper" equations into SymPy Equalities:
>>> from sympy.utilities.iterables import reshape
>>> eqs=reshape(S('''(
x1 = a*s+b*x2 + c
x2 = (m+n) * x3 + d
x3 = l*s + k*x4 + e
x4 = i*s + j*x1 + f)'''.strip().replace(
'\n',',').replace('=',',')), (2,))
>>> eqs = [Eq(*i) for i in eqs]
Then solve the equations (see solve docstring):
>>> solve(eqs, var('x1:5'))
{x4: (a*s + b*d + b*(e + l*s)*(m + n) + c - (-f - i*s)/j)/(-b*k*(m + n) + 1/j),
x3: e + k*(a*s + b*d + b*(e + l*s)*(m + n) + c - (-f - i*s)/j)/(-b*k*(m+ n) + 1/j) +
l*s,
x1: (-f - i*s)/j + (a*s + b*d + b*(e + l*s)*(m + n) + c - (-f - i*s)/j)/(j*(-b*k*(m
+ n) + 1/j)),
x2: d + (m + n)*(e + k*(a*s + b*d + b*(e + l*s)*(m + n) + c - (-f -
i*s)/j)/(-b*k*(m + n) + 1/j) + l*s)}
As noted here, however, this symbolic solution is not valid for all values...it is only valid if none of the coefficients in front of the symbols for which you are solving are zero.
I get a NameError when using SymPy nonlinsolve. After reading a lot of posts I think can be related to type and/or syntax, but I cannot find the exact cause, I just installed sympy a few days ago and my python version is Python 3.5.3 (default, Jan 19 2017, 14:11:04)
[GCC 6.3.0 20170118] on linux
thanks in advance
Vilbjorg
# Python 3 script , using SymPy library, parametrizatrion of the 3-sphere and rotations using quaternion multiplication
# python3 three_sphere.py
from sympy import *
def qmul(x0, x1, x2, x3, y0, y1, y2, y3):
z0 = x0*y0 - x1*y1 - x2*y2 - x3*y3
z1 = x0*y1 + x1*y0 + x2*y3 - x3*y2
z2 = x0*y2 - x1*y3 + x2*y0 + x3*y1
z3 = x0*y3 + x1*y2 - x2*y1 + x3*y0
return z0, z1, z2, z3
r1, s1, t1, r2, s2, t2 = symbols('r1, s1, t1, r2, s2, t2')
a0 = 2*r1/(1 + r1*r1 + s1*s1 + t1*t1)
a1 = 2*s1/(1 + r1*r1 + s1*s1 + t1*t1)
a2 = 2*t1/(1 + r1*r1 + s1*s1 + t1*t1)
a3 = (1 - r1*r1 - s1*s1 - t1*t1)/(1 + r1*r1 + s1*s1 + t1*t1)
b0 = 2*r2/(1 + r2*r2 + s2*s2 + t2*t2)
b1 = 2*s2/(1 + r2*r2 + s2*s2 + t2*t2)
b2 = 2*t2/(1 + r2*r2 + s2*s2 + t2*t2)
b3 = (1 - r2*r2 - s2*s2 - t2*t2)/(1 + r2*r2 + s2*s2 + t2*t2)
c0, c1, c2, c3 = qmul(a0, a1, a2, a3, b0, b1, b2, b3)
c0 = simplify(c0)
c1 = simplify(c1)
c2 = simplify(c2)
c3 = simplify(c3)
print(c0)
print(" ")
print(c1)
print(" ")
print(c2)
print(" ")
print(c3)
print(" ")
print(" ")
r3, s3, t3 = symbols('r3, s3, t3')
q0 = 2*r3/(1 + r3*r3 + s3*s3 + t3*t3)
q1 = 2*s3/(1 + r3*r3 + s3*s3 + t3*t3)
q2 = 2*t3/(1 + r3*r3 + s3*s3 + t3*t3)
q3 = (1 - r3*r3 - s3*s3 - t3*t3)/(1 + r3*r3 + s3*s3 + t3*t3)
#possibly syntax error here which causes NameError ??
soln = nonlinsolve([q0-c0, q1-c1, q2-c2, q3-c3], (r3, s3, t3))
# the idea is to have 4 equations : q0=c0, q1=c1. q2=c2. q3=c3 ; and solve for r3, s3 and t3 in terms of r1, s1, t1, r2, s2, t2
print(soln)
# http://docs.sympy.org/dev/tutorial/solvers.html
You are using a version of SymPy which does not have nonlinsolve. One way to fix this is to update SymPy (versions starting with 1.1 have it). The other is to replace nonlinsolve with solve.
Unfortunately, neither will produce a solution of your system. This is because in general, a system of several algebraic equations has no explicit solution, unless you are really lucky. So neither nonlinsolve nor solve get anywhere with it, despite trying for a long time.
The equation of an ellipse is:
sqrt((x-a1)**2 + (y-b1)**2) + np.sqrt((x-a2)**2 + (y-b2)**2) = c
The focii are (a1, b1) and (a2, b2). c is also known. How do I draw this in python using matplotlib?
Thanks for your help.
You can represent the ellipse parametrically in some variable t. You could look into Wikipedia to see how this can be done, for instance.
In the following code, I've derived the parameters necessary for the parametric form from the parameters that you've supplied.
# Example focii and sum-distance
a1 = 1
b1 = 2
a2 = 5
b2 = 7
c = 9
# Compute ellipse parameters
a = c / 2 # Semimajor axis
x0 = (a1 + a2) / 2 # Center x-value
y0 = (b1 + b2) / 2 # Center y-value
f = np.sqrt((a1 - x0)**2 + (b1 - y0)**2) # Distance from center to focus
b = np.sqrt(a**2 - f**2) # Semiminor axis
phi = np.arctan2((b2 - b1), (a2 - a1)) # Angle betw major axis and x-axis
# Parametric plot in t
resolution = 1000
t = np.linspace(0, 2*np.pi, resolution)
x = x0 + a * np.cos(t) * np.cos(phi) - b * np.sin(t) * np.sin(phi)
y = y0 + a * np.cos(t) * np.sin(phi) + b * np.sin(t) * np.cos(phi)
# Plot ellipse
plt.plot(x, y)
# Show focii
plt.plot(a1, b1, 'bo')
plt.plot(a2, b2, 'bo')
plt.axis('equal')
plt.show()
This gives what you need:
you need 2 lists or arrays of X, Y such that the elements satisfy the ellipse equation
the usual ellipse plotting solutions parameterize the ellipse equation by central (or focal) angle to make the X,Y functions single valued for the angle in 0 to 2pi
I showed a solution in Drawing elliptical orbit in Python (using numpy, matplotlib)
Y as a function of X with a hack to "feel out" the xrange, then piece together the dual Y solutions for each x
just the minimum mods to that code to put in your equation, would fail for a1 = a2
the symbolic solution takes a minute or so runtime
import numpy as np
import matplotlib.pyplot as plt
from sympy import *
# sqrt((x-a1)**2 + (y-b1)**2) + np.sqrt((x-a2)**2 + (y-b2)**2) = c
coeffs = [1, 0, -1, 0, 4]
xs = [coeffs[0], coeffs[2]]
def ysolv(coeffs):
x,y,a1,b1,a2,b2,c = symbols('x y a1 b1 a2 b2 c', real = True)
ellipse = sqrt((x-a1)**2 + (y-b1)**2) + sqrt((x-a2)**2 + (y-b2)**2) - c
y_sols = solve(ellipse, y)
print(*y_sols, sep='\n')
num_coefs = [(a, f) for a, f in (zip([a1,b1,a2,b2,c], coeffs))]
y_solsf0 = y_sols[0].subs(num_coefs)
y_solsf1 = y_sols[1].subs(num_coefs)
print(y_solsf0, '\n', y_solsf1)
f0 = lambdify([x], y_solsf0)
f1 = lambdify([x], y_solsf1)
return f0, f1
f0, f1 = ysolv(coeffs)
y0 = [f0(x) for x in xs]
y1 = [f1(x) for x in xs]
def feeloutXrange(f, midx, endx):
fxs = []
x = midx
while True:
try: f(x)
except:
break
fxs.append(x)
x += (endx - midx)/200
return fxs
midx = (min(xs) + max(xs))/2
xpos = feeloutXrange(f0, midx, max(xs))
xnegs = feeloutXrange(f0, midx, min(xs))
xs_ellipse = xnegs[::-1] + xpos[1:]
y0s = [f0(x) for x in xs_ellipse]
y1s = [f1(x) for x in xs_ellipse]
ys_ellipse = y0s + y1s[::-1] + [y0s[0]] # add y start point to end to close drawing
xs_ellipse = xs_ellipse + xs_ellipse[::-1] + [xs_ellipse[0]] # added x start point
plt.plot(xs_ellipse, ys_ellipse)
plt.show()
(-c*sqrt((a1**2 - 2*a1*a2 + a2**2 + b1**2 - 2*b1*b2 + b2**2 - c**2)*(a1**2 + 2*a1*a2 - 4*a1*x + a2**2 - 4*a2*x + b1**2 - 2*b1*b2 + b2**2 - c**2 + 4*x**2))*(-b1 + b2 + c)*(b1 - b2 + c) + (b1**2 - 2*b1*b2 + b2**2 - c**2)*(-a1**2*b1 + a1**2*b2 + 2*a1*b1*x - 2*a1*b2*x + a2**2*b1 - a2**2*b2 - 2*a2*b1*x + 2*a2*b2*x - b1**3 + b1**2*b2 + b1*b2**2 + b1*c**2 - b2**3 + b2*c**2))/(2*(-b1 + b2 + c)*(b1 - b2 + c)*(b1**2 - 2*b1*b2 + b2**2 - c**2))
(c*sqrt((a1**2 - 2*a1*a2 + a2**2 + b1**2 - 2*b1*b2 + b2**2 - c**2)*(a1**2 + 2*a1*a2 - 4*a1*x + a2**2 - 4*a2*x + b1**2 - 2*b1*b2 + b2**2 - c**2 + 4*x**2))*(-b1 + b2 + c)*(b1 - b2 + c) + (b1**2 - 2*b1*b2 + b2**2 - c**2)*(-a1**2*b1 + a1**2*b2 + 2*a1*b1*x - 2*a1*b2*x + a2**2*b1 - a2**2*b2 - 2*a2*b1*x + 2*a2*b2*x - b1**3 + b1**2*b2 + b1*b2**2 + b1*c**2 - b2**3 + b2*c**2))/(2*(-b1 + b2 + c)*(b1 - b2 + c)*(b1**2 - 2*b1*b2 + b2**2 - c**2))
sqrt(-48*x**2 + 192)/8
-sqrt(-48*x**2 + 192)/8
the other answers used the parametric transform approach
I especially like showing sympy solving the equation for you rather than someone hand solving it
the symbolic expression only needs be found once for a particular Ellipse parameterization, then the symbolic expression can simply be hard coded:
"""
for Ellipse equation:
sqrt((x-a1)**2 + (y-b1)**2) + sqrt((x-a2)**2 + (y-b2)**2) = c
sympy solution to Ellipse equation, only have to run once to get y_sols
symbolic expression to paste into ysolv below
#def symEllipse():
# x,y,a1,b1,a2,b2,c = symbols('x y a1 b1 a2 b2 c', real = True)
# ellipse = sqrt((x-a1)**2 + (y-b1)**2) + sqrt((x-a2)**2 + (y-b2)**2) - c
# y_sols = solve(ellipse, y)
# print(*y_sols, sep='\n')
"""
coeffs = [1, 1, -1, -1, 3]
xs = [coeffs[0], coeffs[2]]
def ysolv(coeffs):
x,y,a1,b1,a2,b2,c = symbols('x y a1 b1 a2 b2 c', real = True)
y_sols = [
(-c*sqrt((a1**2 - 2*a1*a2 + a2**2 + b1**2 - 2*b1*b2 + b2**2 - c**2)*
(a1**2 + 2*a1*a2 - 4*a1*x + a2**2 - 4*a2*x + b1**2 - 2*b1*b2 + b2**2
- c**2 + 4*x**2))*(-b1 + b2 + c)*(b1 - b2 + c) + (b1**2 - 2*b1*b2 +
b2**2 - c**2)*(-a1**2*b1 + a1**2*b2 + 2*a1*b1*x - 2*a1*b2*x +
a2**2*b1 - a2**2*b2 - 2*a2*b1*x + 2*a2*b2*x - b1**3 + b1**2*b2 +
b1*b2**2 + b1*c**2 - b2**3 + b2*c**2))/(2*(-b1 + b2 + c)*
(b1 - b2 + c)*(b1**2 - 2*b1*b2 + b2**2 - c**2)),
(c*sqrt((a1**2 - 2*a1*a2 + a2**2 + b1**2 - 2*b1*b2 + b2**2 - c**2)*
(a1**2 + 2*a1*a2 - 4*a1*x + a2**2 - 4*a2*x + b1**2 - 2*b1*b2 + b2**2
- c**2 + 4*x**2))*(-b1 + b2 + c)*(b1 - b2 + c) + (b1**2 - 2*b1*b2 +
b2**2 - c**2)*(-a1**2*b1 + a1**2*b2 + 2*a1*b1*x - 2*a1*b2*x +
a2**2*b1 - a2**2*b2 - 2*a2*b1*x + 2*a2*b2*x - b1**3 + b1**2*b2 +
b1*b2**2 + b1*c**2 - b2**3 + b2*c**2))/(2*(-b1 + b2 + c)*
(b1 - b2 + c)*(b1**2 - 2*b1*b2 + b2**2 - c**2))
]
num_coefs = [(a, f) for a, f in (zip([a1,b1,a2,b2,c], coeffs))]
y_solsf0 = y_sols[0].subs(num_coefs)
y_solsf1 = y_sols[1].subs(num_coefs)
print(y_solsf0, '\n', y_solsf1)
f0 = lambdify([x], y_solsf0)
f1 = lambdify([x], y_solsf1)
return f0, f1
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.
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
I think that you can expand Matrix expressions. But what you have is not a matrix but the multiplication of two symbolic matrices (Matsymbols). If you turn your expression to a matrix you can get the expansion you wanted. See the extra line below
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 + k2 + k3 + k4)
Matrix(final).expand()