y0 = 1.0
omega = 4.0
h = 0.1
A = matrix([[0.0, 1.0],
[-omega**2, 0.0]])
hA = h * A
A2 = eye(2) + hA + (1.0/2.0)*hA**2
A4 = eye(2) + hA + (1.0/2.0)*hA**2 + (1.0/6.0)*hA**3 + (1.0/24.0)*hA**4
n = int(2*pi/h)
y2, y4 = zeros((n,2)), zeros((n,2))
y2[0,0] = y0
y4[0,0] = y0
t = zeros(n)
for i in range(1,n):
y2[i,:] = A2 @ y2[i-1,:]
y4[i,:] = A4 @ y4[i-1,:]
t[i] = t[i-1] + h
plot(t, cos(omega*t),'--',label="Exact")
plot(t,y2[:,0],label="RK2")
plot(t,y4[:,0],label="RK4")
axis([0, 2*pi,-2,2])
legend(), xlabel("t"), ylabel("$\\theta$");