Those basic equations are solved by the pseudo spectral method.
%config InlineBackend.figure_format = 'retina'
import numpy as np
import matplotlib.pyplot as plt
import scipy.fftpack as sf
from IPython import display
import matplotlib.animation as animation
from tqdm import tqdm
plt.rcParams['font.size'] = 14
plt.rcParams['axes.linewidth'] = 1.5
plt.rcParams['animation.embed_limit']=150
plt.rcParams['animation.html'] = 'jshtml'
def VTE(nx,ny,nt,dt,mu,w,isav):
global dx,KX,KY,KX2,KY2,KXD,KYD
dx=2*np.pi/nx;dy=2*np.pi/ny
### define grids ###
x=np.arange(nx)*dx;y=np.arange(ny)*dy
X,Y=np.meshgrid(x,y)
kx =2*np.pi/(dx*nx)*np.r_[np.arange(nx/2),np.arange(-nx/2,0)]
ky =2*np.pi/(dy*ny)*np.r_[np.arange(ny/2),np.arange(-ny/2,0)]
kxd=np.zeros(nx)
kyd=np.zeros(ny)
kxd[np.abs(kx)<=2/3*np.max(abs(kx))]=1.0 ### 2/3-rule for dealiasing
kyd[np.abs(ky)<=2/3*np.max(abs(ky))]=1.0 ### 2/3-rule for dealiasing
kx2=kx**2; ky2=ky**2
KX,KY =np.meshgrid(kx ,ky )
KX2,KY2=np.meshgrid(kx2,ky2)
KXD,KYD=np.meshgrid(kxd,kyd)
wf=sf.fft2(w)
whst=np.zeros((nt//isav,nx,ny))
for it in tqdm(range(nt)):
### Cranck-Nicholson (only for the dissipation term)###
#wf=1/(1/dt+0.5*mu*(KX2+KY2))*((1/dt+f(wf)+0.5*mu*(KX2+KY2))*wf)
#---double steps with integrating factor method(4th-order Runge-Kutta)---#
wf=np.exp(-mu*(KX2+KY2)*dt)*wf
g1=f(wf )
g2=f(wf+0.5*dt*g1)
g3=f(wf+0.5*dt*g2)
g4=f(wf+ dt*g3)
wf=wf+dt*(g1+2*g2+2*g3+g4)/6
if(it%isav==0):
w =np.real(sf.ifft2(wf))
whst[it//isav,:,:]=w
return locals()
def f(wf):
phif=wf/(KX2+KY2+1e-10); phif[0,0]=0.0
vxf = 1j*KY*phif; vx=np.real(sf.ifft2(vxf*KXD*KYD))
vyf =-1j*KX*phif; vy=np.real(sf.ifft2(vyf*KXD*KYD))
wxf = 1j*KX*wf ; wx=np.real(sf.ifft2(wxf*KXD*KYD))
wyf = 1j*KY*wf ; wy=np.real(sf.ifft2(wyf*KXD*KYD))
advw =vx*wx+vy*wy
advwf=sf.fft2(advw)
return -advwf
nx=512; ny=512; nt=1200; isav=50
mu=0.0002; dt=0.1
w=np.random.randn(nx,ny)*5
dat=VTE(nx,ny,nt,dt,mu,w,isav)
locals().update(dat)
def update_anim(it):
fig.clf()
ax1 = fig.add_subplot(221)
ax2 = fig.add_subplot(222)
ax3 = fig.add_subplot(223)
ax4 = fig.add_subplot(224)
for ax in (ax1, ax2, ax3, ax4):
ax.clear()
wf =sf.fft2(whst[it,:,:])
nabla2w=np.real(sf.ifft2(-(KX2+KY2)*wf))
phif=wf/(KX2+KY2+1e-10); phif[0,0]=0.0
phi =np.real(sf.ifft2(phif))
phixyf=-KX*KY*phif; phixy=np.real(sf.ifft2(phixyf*KXD*KYD))
phix2f=-KX2*phif ; phix2=np.real(sf.ifft2(phix2f*KXD*KYD))
phiy2f=-KY2*phif ; phiy2=np.real(sf.ifft2(phiy2f*KXD*KYD))
vxf = 1j*KY*phif; vx=np.real(sf.ifft2(vxf*KXD*KYD))
vyf =-1j*KX*phif; vy=np.real(sf.ifft2(vyf*KXD*KYD))
wxf = 1j*KX*wf ; wx=np.real(sf.ifft2(wxf*KXD*KYD))
wyf = 1j*KY*wf ; wy=np.real(sf.ifft2(wyf*KXD*KYD))
advw =-(vx*wx+vy*wy)
p8=2*phixy**2+2*phix2*phiy2
im1=ax1.imshow(whst[it,:,:] ,aspect='auto',origin='lower',cmap='jet');ax1.axis('off');fig.colorbar(im1, ax=ax1);ax1.set_title(r'$\omega$')
im2=ax2.imshow(p8 ,aspect='auto',origin='lower',cmap='jet');ax2.axis('off');fig.colorbar(im2, ax=ax2);ax2.set_title(r'$R_{\nu}^{-1}\nabla^2\omega$')
im3=ax3.imshow(advw ,aspect='auto',origin='lower',cmap='jet');ax3.axis('off');fig.colorbar(im3, ax=ax3);ax3.set_title(r'$-({\bf v}\cdot\nabla)\omega$')
im4=ax4.imshow(phi ,aspect='auto',origin='lower',cmap='jet');ax4.axis('off');fig.colorbar(im4, ax=ax4);ax4.set_title(r'$\phi\ (\omega\equiv\nabla^2\phi)$')
fig=plt.figure(figsize=(15,12))
anim=animation.FuncAnimation(fig,update_anim,frames=nt//isav)
plt.close()
anim