y0 = 1.0
omega = 4.0
h = 0.1
A = array([[0.0, 1.0],
[-omega**2, 0.0]])
Afe = eye(2) + h * A
Abe = inv(eye(2) - h * A)
At1 = eye(2) - 0.5 * h * A
At2 = eye(2) + 0.5 * h * A
Atr = inv(At1) @ At2
n = int(2*pi/h)
yfe, ybe, ytr = zeros((n,2)), zeros((n,2)), zeros((n,2))
yfe[0,0] = y0
ybe[0,0] = y0
ytr[0,0] = y0
t = zeros(n)
for i in range(1,n):
yfe[i,:] = Afe @ yfe[i-1,:]
ybe[i,:] = Abe @ ybe[i-1,:]
ytr[i,:] = Atr @ ytr[i-1,:]
t[i] = t[i-1] + h
plot(t, cos(omega*t),'--',label="Exact")
plot(t,yfe[:,0],label="FE")
plot(t,ybe[:,0],label="BE")
plot(t,ytr[:,0],label="TR")
axis([0, 2*pi,-2,2])
legend(), xlabel("t"), ylabel("$\\theta$");