import numpy as np def time_integrate(K,N): T = 0.2 tau = T/N x_K = np.linspace(-np.pi,np.pi,K,endpoint=False) U0 = np.exp(-10*x_K**2+10*1j*x_K) U = 1j*np.zeros((K,N+1)) U[:,0] = U0 D_K = -np.diag(np.arange(-K/2,K/2)**2) V_K = np.diag(100*(1-np.cos(x_K))) for n in range(N): DK_FK_f = D_K.dot(np.fft.fftshift(np.fft.fft(U[:,n]))) lap_approx_Un = np.fft.ifft(np.fft.ifftshift(DK_FK_f)) U[:,n+1] = U[:,n] +tau*1j*lap_approx_Un -tau*1j*V_K.dot(U[:,n]) return U def make_images(K,U): plt.semilogy(1.0/np.sqrt(K)*np.linalg.norm(U,axis=0)) plt.ylim([10**(-1),10**(0)]) plt.savefig('norms_over_time') plt.close() plt.imshow(np.abs(U)**2,interpolation='nearest',aspect='auto') plt.colorbar() plt.savefig('solution') import matplotlib.pyplot as plt K = 32 N = 10000 U = time_integrate(K,N) print("Finished") #plt.semilogy(np.linalg.norm(U,axis = 0)) #plt.savefig("norms") make_images(K,U)