def inicheck(dirname):
import numpy as np
import os
from pylab import *
import fortranfile
from math import erf
import pickle
ion()
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
print L, beta_dim,f0_dim, h1
U=beta_dim*L**2/beta1
gprime=f0_dim**2*L**2/h1/fr1
h2=f0_dim**2*L**2/gprime/fr2
#are jsi and jso different? jsdiff=0 if they are not, 1 if they are.
# jsdiff=1
x=linspace(xmin,xmax,n)
y=linspace(ymin,ymax,m)
from makefields import getinitPsi,gethb,gettracer
[PSI1,PSI2]=getinitPsi(m,n,y,gs1,g1,dw1,d1,as2,a2,movegs,movedw,nspo,mspo)
hb=gethb(cm,cz,ymax,movedw,hw,hbmax,a2,n,m)
ag=gettracer(ndye,athick,ymax,movedw,a2,n,m)
ymi=3*m/4
xmi=0
xma=10
figure()
plot(y[ymi:],(PSI2[0,ymi:]/PSI2.max()+1)/2,'bx-',label='Psi')
plot(y[ymi:],hb[0,ymi:]/hbmax,'rx-',label='H')
#plot(y[ymi:],ag[:,ymi:].T,'g')
#plot(y[ymi:],ag[0,ymi:],'g',label='tracer')
xlabel('y position')
legend()
title('Normalized $\psi$ and topography')
grid('on')
savefig('../plots/'+dirname+'/'+dirname[7:]+'_initopopsi.png')
#get a grasp on the transport
from operators import grady
dy=(ymax-ymin)/real(m-1)
t4dy=1.0/(8.0*dy)
u1=grady(PSI1,t4dy,n,m)
u2=grady(PSI2,t4dy,n,m)
u1=-u1*U
u2=-u2*U
ymi=0
mind=1
figure()
plot(y[ymi:],u1[mind,ymi:],'b',label='Layer 1')
plot(y[ymi:],u2[mind,ymi:],'r',label='Layer 2')
xlabel('y position')
#legend()
title('velocity')
grid('on')
savefig('../plots/'+dirname+'/'+dirname[7:]+'_inivel.png')
#layer interface height
hm=f0_dim/gprime*(PSI2-PSI1)*U*L
hbp=hb*h2*U/f0_dim/L
hm1=h1-hm
hm2=h2+hm-hbp
tm1=u1[mind,:]*hm1[mind,:]*L*(xmax-xmin)/n/1e6
tm2=u2[mind,:]*hm2[mind,:]*L*(xmax-xmin)/n/1e6
ctm1=cumsum(tm1)
ctm2=cumsum(tm2)
figure()
ax=subplot(221)
plot(y/1000,tm1)
ylabel('Transport per grid point (Sv)')
title('Upper Layer, Total= '+str(around(ctm1[-1],decimals=1))+' Sv')
ax.set_xticklabels([])
grid(True)
ax=subplot(222)
plot(y/1000,tm2)
title('Lower Layer, Total= '+str(around(ctm2[-1],decimals=1))+' Sv')
ax.set_xticklabels([])
grid(True)
ax=subplot(223)
ylabel('Cumulative transport (Sv)')
plot(y/1000,ctm1)
xlabel('y (1000 km)')
grid(True)
ax=subplot(224)
plot(y/1000,ctm2)
xlabel('y (1000 km)')
grid(True)
savefig('../plots/'+dirname+'/'+dirname[7:]+'_initrans.png')
figure()
plot(y[ymi:],u2[mind,ymi:]/u2[mind,ymi:].min(),'r',label='velocity')
plot(y[ymi:],hm2[mind,ymi:]/hm2[mind,ymi:].max(),'b',label='layer thickness')
plot(y[ymi:],tm2[ymi:]/tm2[ymi:].min(),'m',label='transport')
xlabel('y position')
#legend()
#title('velocity')
grid('on')
savefig('../plots/'+dirname+'/'+dirname[7:]+'_inicomp.png')
def pvbudload(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
import os
#ion()
from makefields import getrspo,getinitPsi,gethb
from operators import delsq,gradx,grady
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
U=beta_dim*L**2/beta1
gprime=f0_dim**2*L**2/h1/fr1
h2=f0_dim**2*L**2/gprime/fr2
f0_nd=f0_dim*L/U
#define x and y (m)
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
dx=(xmax-xmin)/real(n-1)
dy=(ymax-ymin)/real(m-1)
dx2=dx**2
t4dx=1.0/(8.0*dx)
t4dy=1.0/(8.0*dy)
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
predir='/data/ilebras/twolayer-jets-climate/'
rundir=predir+'testing/'+dirname+'/'
# --------load fields-------
from loadfuncs import load1
#mean pv fluxes in layer 1
mf1=load1('mf1.dat',rundir,n,m)
vm1=load1('vm1.dat',rundir,n,m)
um1=load1('um1.dat',rundir,n,m)
mfi=load1('mfi.dat',rundir,n,m)
#mean pv fluxes in layer 2
mf2=load1('mf2.dat',rundir,n,m)
vm2=load1('vm2.dat',rundir,n,m)
um2=load1('um2.dat',rundir,n,m)
betav1=beta1*vm1
mthi1=fr1*mfi
mpvf1=mf1+betav1+mthi1
# topographic term
hb=gethb(cm,cz,ymax,movedw,hw,hbmax,a2,n,m)
hby=grady(hb,t4dy,n,m)
topo=vm2*hby
# bottom drag
mz2=load1('z2m.dat',rundir,n,m)
bdrag=mz2/rdecayw
betav2=beta2*vm2
mthi2=-fr2*mfi
mpvf2=mf2+betav2+mthi2+topo+bdrag
#eddy pv fluxes in layer 1
ef1=load1('ef1.dat',rundir,n,m)
efi=load1('efi.dat',rundir,n,m)
#eddy pv fluxes in layer 2
ef2=load1('ef2.dat',rundir,n,m)
ef2[abs(ef2)>10**13]=0.0
ef2[isnan(ef2)]=0.0
ethi1=fr1*efi
epvf1=ef1+ethi1
ethi2=-fr2*efi
epvf2=ef2+ethi2
#viscosity and sponge drag
mz1=load1('z1m.dat',rundir,n,m)
visc1=-akap*delsq(mz1,dx2,n,m)
visc2=-akap*delsq(mz2,dx2,n,m)
#sponge drag
rspo=getrspo(nspo,mspo,rdecay,rdecayo,n,m,dt)
[p1init,p2init]=getinitPsi(m,n,y/L,gs1,g1,dw1,d1,as2,a2,movegs,movedw,nspo,mspo)
z1init=delsq(p1init,dx2,n,m)
z2init=delsq(p2init,dx2,n,m)
p1xinit=gradx(p1init,t4dx,n,m)
p1yinit=grady(p1init,t4dy,n,m)
p2xinit=gradx(p2init,t4dx,n,m)
p2yinit=grady(p2init,t4dy,n,m)
rspox=gradx(rspo,t4dx,n,m)
rspoy=grady(rspo,t4dy,n,m)
sponge1_del=rspo*(mz1-z1init)
sponge1_cross=rspox*(vm1-p1xinit)-rspoy*(um1+p1yinit)
sponge1=sponge1_del+sponge1_cross
sponge2_del=rspo*(mz2-z2init)
sponge2_cross=rspox*(vm2-p2xinit)-rspoy*(um2+p2yinit)
sponge2=sponge2_del+sponge2_cross
resid1=mpvf1+epvf1+visc1+sponge1
resid2=mpvf2+epvf2+visc2+sponge2
#return p1init,p2init,z1init,z2init,p1xinit,p1yinit,p2xinit,p2yinit,hb,hby
pickle.dump([mf1,betav1,mthi1,mpvf1,ef1,ethi1,epvf1,visc1,sponge1,resid1,mf2,betav2,mthi2,mpvf2,ef2,ethi2,epvf2,visc2,sponge2,topo,bdrag,resid2,hb,hby], open( predir+'pickles/pvbud/'+dirname+'_pvbud' , 'w'))
def pvfluxload(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
import os
# ion()
# %matplotlib inline
from makefields import getrspo,getinitPsi,gethb
from operators import delsq,gradx,grady
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
U=beta_dim*L**2/beta1
gprime=f0_dim**2*L**2/h1/fr1
h2=f0_dim**2*L**2/gprime/fr2
f0_nd=f0_dim*L/U
#define x and y (m)
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
dx=(xmax-xmin)/real(n-1)
dy=(ymax-ymin)/real(m-1)
dx2=dx**2
t4dx=1.0/(8.0*dx)
t4dy=1.0/(8.0*dy)
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
predir='/data/ilebras/twolayer-jets-climate/'
rundir=predir+'testing/'+dirname+'/'
# --------load fields-------
from loadfuncs import load1
#mean pv fluxes
um1=load1('um1.dat',rundir,n,m)
um2=load1('um2.dat',rundir,n,m)
vm1=load1('vm1.dat',rundir,n,m)
vm2=load1('vm2.dat',rundir,n,m)
zm1=load1('z1m.dat',rundir,n,m)
zm2=load1('z2m.dat',rundir,n,m)
qm1=load1('qm1.dat',rundir,n,m)
qm2=load1('qm2.dat',rundir,n,m)
pm1=load1('pm1.dat',rundir,n,m)
pm2=load1('pm2.dat',rundir,n,m)
hb=gethb(cm,cz,ymax,movedw,hw,hbmax,a2,n,m)
# qm1=zm1+betay+fr1*(pm2-pm1)#+f0_nd
qm2_check=zm2+betay+fr2*(pm1-pm2)+hb#+f0_nd
# figure()
# plot(qm2.T);
# plot(qm2_check.T);
# figure()
# subplot(211)
# plot(um2.T)
# subplot(212)
# plot(vm2.T)
uqm1=um1*qm1
uqm2=um2*qm2
vqm1=vm1*qm1
vqm2=vm2*qm2
# #calculate divergence
# duqm1x=gradx(uqm1,t4dx,n,m)
# duqm2x=gradx(uqm2,t4dx,n,m)
# dvqm1y=grady(vqm1,t4dx,n,m)
# dvqm2y=grady(vqm2,t4dx,n,m)
# divuq1m=duqm1x+dvqm1y
# divuq2m=duqm2x+dvqm2y
#eddy pv fluxes
uz1=-(load1('uz1.dat',rundir,n,m)+2*um1*zm1)
uz2=-(load1('uz2.dat',rundir,n,m)+2*um2*zm2)
vz1=load1('vz1.dat',rundir,n,m)
vz2=load1('vz2.dat',rundir,n,m)
#need to correct for a missing minus sign in stats2_pvbud.F
u1p1=-(load1('u1p1.da',rundir,n,m)+2*um1*pm1)
u1p2=-(load1('u1p2.da',rundir,n,m)+2*um1*pm2)
u2p1=-(load1('u2p1.da',rundir,n,m)+2*um2*pm1)
u2p2=-(load1('u2p2.da',rundir,n,m)+2*um2*pm2)
v1p1=load1('v1p1.da',rundir,n,m)
v1p2=load1('v1p2.da',rundir,n,m)
v2p1=load1('v2p1.da',rundir,n,m)
v2p2=load1('v2p2.da',rundir,n,m)
uthi1=fr1*(u1p2-u1p1)
uthi2=fr2*(u2p1-u2p2)
vthi1=fr1*(v1p2-v1p1)
vthi2=fr2*(v2p1-v2p2)
uqb1=uz1+uthi1
uqb2=uz2+uthi2
vqb1=vz1+vthi1
vqb2=vz2+vthi2
dqy1=grady(qm1,t4dx,n,m)
dqy2=grady(qm2,t4dx,n,m)
uz={}
uz['1']=uz1
uz['2']=uz2
vz={}
vz['1']=vz1
vz['2']=vz2
ub={}
ub['1']=uqb1
ub['2']=uqb2
vb={}
vb['1']=vqb1
vb['2']=vqb2
up={}
up['11']=u1p1
up['12']=u1p2
up['21']=u2p1
up['22']=u2p2
vp={}
vp['11']=v1p1
vp['12']=v1p2
vp['21']=v2p1
vp['22']=v2p2
uqm={}
uqm['1']=uqm1
uqm['2']=uqm2
vqm={}
vqm['1']=vqm1
vqm['2']=vqm2
pickle.dump([uz,vz,ub,vb,up,vp,uqm,vqm,dqy1,dqy2],open(predir+'pickles/pvfluxmv/'+dirname+'_pvflmv.p', 'w'))
def qls(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
import os
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
U=beta_dim*L**2/beta1
gprime=f0_dim**2*L**2/h1/fr1
h2=f0_dim**2*L**2/gprime/fr2
f0_nd=f0_dim*L/U
#define x and y (m)
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
predir='/data/ilebras/twolayer-jets-climate/'
rundir=predir+'testing/'+dirname+'/'
#--------move on to load psi--------------
psi={}
pload=fortranfile.FortranFile(rundir+'pm1.dat')
pvec=pload.readReals()
pm1=pvec.reshape(n,m,order='F')
psi['m1']=pm1*U*L
pload=fortranfile.FortranFile(rundir+'pm2.dat')
pvec=pload.readReals()
pm2=pvec.reshape(n,m,order='F')
psi['m2']=pm2*U*L
# pload=fortranfile.FortranFile(rundir+'psidat/p1i_000')
# pvec=pload.readReals()
# p1i=pvec.reshape(n,m,order='F')
# psi['i1']=p1i*U*L
# pload=fortranfile.FortranFile(rundir+'psidat/p2i_000')
# pvec=pload.readReals()
# p2i=pvec.reshape(n,m,order='F')
# psi['i2']=p2i*U*L
# #--------and then load q (zeta)--------------
q={}
pload=fortranfile.FortranFile(rundir+'z1m.dat')
pvec=pload.readReals()
qm1=pvec.reshape(n,m,order='F')
q['m1']=qm1*U/L
pload=fortranfile.FortranFile(rundir+'z2m.dat')
pvec=pload.readReals()
qm2=pvec.reshape(n,m,order='F')
q['m2']=qm2*U/L
# pload=fortranfile.FortranFile(rundir+'psidat/q1i_000')
# pvec=pload.readReals()
# q1i=pvec.reshape(n,m,order='F')
# q['i1']=q1i*U/L
# pload=fortranfile.FortranFile(rundir+'psidat/q2i_000')
# pvec=pload.readReals()
# q2i=pvec.reshape(n,m,order='F')
# q['i2']=q2i*U/L
#--------if there is bottom topography, plot it-----
from makefields import gethb
hb=gethb(cm,cz,ymax,movedw,hw,hbmax,a2,n,m)
hbp=hb*h2*U/f0_dim/L
#--------and then plot qfull= q +fr(psi-psi, depending) + betay + hb--------
qfull={}
qfull['m1']=q['m1']*L/U+betay+fr1*(psi['m2']-psi['m1'])/(U*L)#+f0_nd
qfull['m2']=q['m2']*L/U+betay+fr2*(psi['m1']-psi['m2'])/(U*L)+hb#+f0_nd
# qfull['i1']=q['i1']*L/U+betay+fr1*(psi['i2']-psi['i1'])/(U*L)#+f0_nd
# qfull['i2']=q['i2']*L/U+betay+fr2*(psi['i1']-psi['i2'])/(U*L)+hb#+f0_nd
#--------calculate a few more things leading up to transport--------
u={}
v={}
from operators import gradx, grady
dx=(xmax-xmin)/real(n-1)
dy=(ymax-ymin)/real(m-1)
dx2=dx**2
t4dx=1.0/(8.0*dx)
t4dy=1.0/(8.0*dy)
# u['m1']=-diff(psi['m1'],n=1,axis=1)/diff(y,n=1)
# u['m2']=-diff(psi['m2'],n=1,axis=1)/diff(y,n=1)
# u['i1']=-diff(psi['i1'],n=1,axis=1)/diff(y,n=1)
# u['i2']=-diff(psi['i2'],n=1,axis=1)/diff(y,n=1)
# v['m1']=(diff(psi['m1'],n=1,axis=0).T/diff(x,n=1)).T
# v['m2']=(diff(psi['m2'],n=1,axis=0).T/diff(x,n=1)).T
# v['i1']=(diff(psi['i1'],n=1,axis=0).T/diff(x,n=1)).T
# v['i2']=(diff(psi['i2'],n=1,axis=0).T/diff(x,n=1)).T
u['m1']=-grady(psi['m1']/(U*L),t4dy,n,m)*U
u['m2']=-grady(psi['m2']/(U*L),t4dy,n,m)*U
# u['i1']=-grady(psi['i1']/(U*L),t4dy,n,m)*U
# u['i2']=-grady(psi['i2']/(U*L),t4dy,n,m)*U
v['m1']=gradx(psi['m1']/(U*L),t4dx,n,m)*U
v['m2']=gradx(psi['m2']/(U*L),t4dx,n,m)*U
# v['i1']=gradx(psi['i1']/(U*L),t4dx,n,m)*U
# v['i2']=gradx(psi['i2']/(U*L),t4dx,n,m)*U
#note that u,v is nondimensional here.
#calculate layer depth fluctuations from psi
hm=f0_dim/gprime*(psi['m2']-psi['m1'])
#hi=f0_dim/gprime*(psi['i2']-psi['i1'])
h={}
h['m1']=h1-hm
h['m2']=h2+hm-hbp
# h['i1']=h1-hi
# h['i2']=h2+hi-hbp
if not os.path.exists(predir+'pickles/quickfields/'+dirname[:6]):
os.mkdir(predir+'pickles/quickfields/'+dirname[:6])
pickle.dump([psi,q,qfull,u,v,h,hbp,x,y], open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'w' ) )
def movie_pv(dirname,iors='s'):
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
import matplotlib.cm as cmapo
ion()
predir='/data/ilebras/twolayer-jets-climate/'
pathint='testing/'+dirname+'/psidat/'
layer='2'
if not os.path.exists('../plots/'+dirname+'/q'+layer+'zoom'+iors+'mov'):
#os.mkdir('../plots/'+dirname)
os.mkdir('../plots/'+dirname+'/q'+layer+'zoom'+iors+'mov')
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
U=beta_dim*L**2/beta1
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
f0_nd=f0_dim*L/U
ylow=200
if iors=='s':
secspace=0.02/dt
tmaxi=ksnap
elif iors=='i':
secspace=0.01
tmaxi=ktmax
time=arange(ksec,tmaxi,1/secspace)*L/U/60**2/24/365.25
tdim=len(time)
#pause time when watching psi movie
ptime=0.0001
from makefields import gethb
hb=gethb(cm,cz,ymax,movedw,hw,hbmax,a2,n,m)
pload=fortranfile.FortranFile(predir+pathint+'q2'+iors+'_000')
pvec=pload.readReals()
q2=pvec.reshape(n,m,order='F')
pload=fortranfile.FortranFile(predir+pathint+'q1'+iors+'_000')
pvec=pload.readReals()
q1=pvec.reshape(n,m,order='F')
pload=fortranfile.FortranFile(predir+pathint+'p2'+iors+'_000')
pvec=pload.readReals()
p2=pvec.reshape(n,m,order='F')
pload=fortranfile.FortranFile(predir+pathint+'p1'+iors+'_000')
pvec=pload.readReals()
p1=pvec.reshape(n,m,order='F')
#qfull0=q2+betay+fr2*(p1-p2)+hb#+f0_nd
qfull0=q1+betay+fr1*(p2-p1)+hb
#set the range for the colorbar to be greater than zero and to the max
#cmin1=qfull0.min()
#cmax1=qfull0.max()
cmin1=-0.5
cmax1=4
#cmax1=0.5
setc1=linspace(cmin1,int(cmax1),501)
xmini=mspo
xmaxi=-3*mspo/2
ydw=int(m-(ymax-movedw)*m/2/ymax)
ytopo=int(m-(ymax-movedw-hw/20)*m/2/ymax-hw)
ytopo2=int(m-(ymax-movedw-hw/20)*m/2/ymax)
fs=18
figure(figsize=(20,6))
for iter in range(0,999):
print iter
clf()
q2load=fortranfile.FortranFile(predir+pathint+'q2'+iors+'_'+str(iter).zfill(3))
q2vec=q2load.readReals()
q2=q2vec.reshape(n,m,order='F')
q1load=fortranfile.FortranFile(predir+pathint+'q1'+iors+'_'+str(iter).zfill(3))
q1vec=q1load.readReals()
q1=q1vec.reshape(n,m,order='F')
p1load=fortranfile.FortranFile(predir+pathint+'p1'+iors+'_'+str(iter).zfill(3))
p1vec=p1load.readReals()
p1=p1vec.reshape(n,m,order='F')
p2load=fortranfile.FortranFile(predir+pathint+'p2'+iors+'_'+str(iter).zfill(3))
p2vec=p2load.readReals()
p2=p2vec.reshape(n,m,order='F')
if layer=='1':
qfull=q1+betay+fr1*(p2-p1)
elif layer=='2':
qfull=q2+betay+fr2*(p1-p2)+hb#+f0_nd
qfull[qfull<=cmin1]=cmin1
qfull[qfull>=cmax1]=cmax1
cla()
contourf(x[xmini:xmaxi]/1e3,y[ylow:]/1e3,qfull[xmini:xmaxi,ylow:].T,setc1,cmap='RdBu')
cbar=colorbar()
cbar.set_label('PV (nondim)',fontsize=fs)
# plot(x[mspo/2]/1000*ones(shape(y)),y/1000,'k')
# plot(x[-mspo]/1000*ones(shape(y)),y/1000,'k')
plot(x/1000,y[ydw]/1000*ones(shape(x)),'k')
plot(x/1000,y[ytopo]/1000*ones(shape(x)),'k')
plot(x/1000,y[ytopo2]/1000*ones(shape(x)),'k')
xlim([x[xmini]/1e3, x[xmaxi-1]/1e3])
ylim([y[ylow:].min()/1e3, y[ylow:].max()/1e3])
xlabel('x position (km)',fontsize=fs)
ylabel('y position (km)',fontsize=fs)
if layer=='1':
title('Upper Layer',fontsize=fs+2)
elif layer=='2':
title('Lower Layer',fontsize=fs+2)
#load instantaneous PV to plot on there too
savefig(predir+'plots/'+dirname+'/q'+layer+'zoom'+iors+'mov/'+str(dirname[11:])+'_q'+layer+'zoom_'+str(int(iter)).zfill(3)+'.png')