import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
zeta_list = np.array([0, 0.2, 0.4, 0.6, 0.7, 0.8, 0.99])
omega = 1
K = len(zeta_list)
Y = []
for zeta in zeta_list:
aa = np.array([[0, 1],[-omega**2, -2*zeta*omega]])
bb = np.array([[0], [omega**2]])
cc = np.array([[1, 0]])
dd = 0
G = signal.StateSpace(aa,bb,cc,dd)
t, yout, xout = signal.lsim(G, U=0, T=np.linspace(0,10,1000), X0=[1,0])
Y.append(yout)
plt.figure(figsize=(14,9), dpi=100)
for k in range(K):
zeta = zeta_list[k]
plt.plot(t, Y[k], label=rf'$\zeta={zeta:3.1f}$', linewidth=2, color=plt.cm.tab20(0.1*k))
plt.plot(t, np.exp(-zeta*omega*t)*np.cos(np.sqrt(1-zeta**2)*omega*t-np.arctan2(zeta,np.sqrt(1-zeta**2)))/np.sqrt(1-zeta**2), \
linewidth=2, linestyle='--', alpha=0.5, color=plt.cm.tab20(0.1*k))
UB = np.exp(-zeta_list[k]*omega*t)
plt.fill_between(t, -UB, UB, alpha=0.05)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x$')
plt.title(r'Solutions of $\ddot{x} + 2\zeta\omega\dot{x} + \omega^2 x = 0$ with $\omega=1$ and $(x_0, \dot{x}_0)=(1,0)$')
plt.legend()
plt.grid()
plt.show()