import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
%matplotlib inline
a = 1
lmbda = 0.6/a
dxs= np.array([.1, .025])
dts = dxs*lmbda
ks = np.arange(1,16)
omega_R_FTBS = np.zeros((len(dts),len(ks)))
omega_I_FTBS = np.zeros((len(dts),len(ks)))
omega_R_LW = np.zeros((len(dts),len(ks)))
omega_I_LW = np.zeros((len(dts),len(ks)))
for i in range(len(dts)):
for k in ks:
dt = dts[i]
dx = dxs[i]
omega_R_FTBS[i,k-1] = np.arctan(a*lmbda*np.sin(k*dx)/(1 - a*lmbda + a*lmbda*np.cos(k*dx)))/dt
omega_I_FTBS[i,k-1] = np.log(a*lmbda*np.sin(k*dx)/np.sin(omega_R_FTBS[i,k-1]*dt))/dt
omega_R_LW[i,k-1] = np.arctan(a*lmbda*np.sin(k*dx)/(1 - a**2*lmbda**2 + a**2*lmbda**2*np.cos(k*dx)))/dt
omega_I_LW[i,k-1] = np.log(a*lmbda*np.sin(k*dx)/np.sin(omega_R_LW[i,k-1]*dt))/dt
plt.plot(ks, omega_R_FTBS[0,:], lw=3, label='FTBS')
plt.plot(ks, omega_R_LW[0,:], lw=3, label='LW')
plt.plot(ks, a*ks, linewidth=2, color='black', label='exact')
plt.legend(frameon=False, loc='upper left')
plt.plot( ks, omega_I_FTBS[0,:], lw=3, label='FTBS')
plt.plot( ks, omega_I_LW[0,:], lw=3, label='LW')
plt.legend(frameon=False, loc='upper left')