連成振動をグラフ化してみた

matplotlibを使って昨日書いた連成振動をグラフ化してみた。

#!/usr/bin/python
# -*- coding:utf-8 -*-

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

#-- initialization
m = 1.0
k = 1.0
kappa = 1.0
dt = 0.01
omega1 = (k + kappa) / m
omega2 = kappa / m
stepNum = 10000
freq = 10

x1 = 0.0
x2 = 0.0
v1 = 1.0
v2 = -0.5

mu = 0.0
sigma = 1.0

#-- state equation
x = np.matrix([[x1],[x2],[v1],[v2]])
A = np.matrix([[0.0,0.0,1.0,0.0],
             [0.0,0.0,0.0,1.0],
             [-omega1,omega2,0.0,0.0],
             [omega2,-omega1,0.0,0.0]])

#-- observation equation
H = np.matrix([[1.0,0.0,0.0,0.0],
             [0.0,1.0,0.0,0.0]])
z = H * x

#-- list for data plotting
xaxis = []
z0 = []
z1 = []

#-- Runge-Kutta 4th order Method
for t in range(stepNum):
    if(t % freq == 0):
        tv = 0.05 * np.matrix(stats.norm.rvs(0,1,size = 2))
        v = tv.transpose()
        z = H * x + v
        xaxis.append(t*dt)
        z0.append(float(z[0]))
        z1.append(float(z[1]))
    
    k1 = dt * A * x
    k2 = dt * A * (x + k1/2.0)
    k3 = dt * A * (x + k2/2.0)
    k4 = dt * A * (x + k3)
    x = x + (k1 + 2.0*k2 + 2.0*k3 + k4)/6.0  

plt.xlabel('time')
plt.ylabel('z')
plt.plot(xaxis,z0)
plt.plot(xaxis,z1)
plt.show()

軸とか色とかキャンバスサイズとかは適当。
一応連成振動(観測誤差入り)しているのがわかる。