def loadTEM(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
import os
from pf2d_ztitvar import pf2d_ztitvar
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+'/'
#--------load eddy forcing terms--------------
uq_TEM={}
pload=fortranfile.FortranFile(rundir+'Tr1.dat')
pvec=pload.readReals()
ppp=pvec.reshape(n,m,order='F')
uq_TEM['m1']=ppp
pload=fortranfile.FortranFile(rundir+'Tr2.dat')
pvec=pload.readReals()
ppp=pvec.reshape(n,m,order='F')
uq_TEM['m2']=ppp
pload=fortranfile.FortranFile(rundir+'TM1.dat')
pvec=pload.readReals()
ppp=pvec.reshape(n,m,order='F')
uq_TEM['i1']=ppp
pload=fortranfile.FortranFile(rundir+'TM2.dat')
pvec=pload.readReals()
ppp=pvec.reshape(n,m,order='F')
uq_TEM['i2']=ppp
#--------plot eddy forcing terms--------------
xlow=200
xhigh=-3
ylow=250
pf2d_ztitvar(uq_TEM,'uqTEM_fromsaved_no',x,y,dirname,(18,6),xlow,xhigh,ylow,'relative vorticity portion','total eddy forcing',1)
def tractseries(finaldir):
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
predir = "/data/ilebras/twolayer-jets-climate/"
from readjet2h import readjet2h
vardict = readjet2h(finaldir)
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
x = arange(xmin, xmax, (xmax - xmin) / n) * L
y = arange(ymin, ymax, (ymax - ymin) / m) * L
secspace = 0.02 / dt
time = arange(ksec, ksnap, 1 / secspace) * L / U / 60 ** 2 / 24 / 365.25
tdim = len(time)
amean1 = ones(tdim)
amin1 = zeros(tdim)
amax1 = zeros(tdim)
asum1 = zeros(tdim)
for iter in range(tdim):
# load some age tracers
pload = fortranfile.FortranFile("../testing/" + finaldir + "/psidat/a1s_" + str(int(iter)).zfill(3))
pvec = pload.readReals()
age = pvec.reshape(n, m, order="F")
amean1[iter] = np.mean(age)
amin1[iter] = age.min()
amax1[iter] = age.max()
asum1[iter] = sum(age[:, :])
amean2 = ones(tdim)
amin2 = zeros(tdim)
amax2 = zeros(tdim)
asum2 = zeros(tdim)
for iter in range(tdim):
# load some age tracers
pload = fortranfile.FortranFile("../testing/" + finaldir + "/psidat/a2s_" + str(int(iter)).zfill(3))
pvec = pload.readReals()
age = pvec.reshape(n, m, order="F")
amean2[iter] = np.mean(age)
amin2[iter] = age.min()
amax2[iter] = age.max()
asum2[iter] = sum(age[:, :])
figure()
subplot(121)
xlabel("time (years)")
plot(time, amean1, color="b")
plot(time, amin1, color="b")
plot(time, amax1, color="b")
plot(time, amean2, color="r")
plot(time, amin2, color="r")
plot(time, amax2, color="r")
title("Tracer mean, min and max")
grid(True)
subplot(122)
plot(time, asum1, "b")
plot(time, asum2, "r")
title("Tracer sum")
legend(["layer 1", "layer 2"])
xlabel("time (years)")
grid(True)
savefig(predir + "plots/" + finaldir + "_tracertseries.png")
def plotu(dirname):
import pickle
from pylab import *
import math
from uplotfunc import upf
ion()
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
predir='/data/ilebras/twolayer-jets-climate/'
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
U=beta_dim*L**2/beta1
# xmid=x[:-1]+diff(x)/2
# ymid=y[:-1]+diff(y)/2
[xmmat,ymmat]=meshgrid(x,y)
# um1p=u['m1'][:-1,:]+diff(u['m1'],axis=0)/2
# vm1p=v['m1'][:,:-1]+diff(v['m1'],axis=1)/2
# ui1p=u['i1'][:-1,:]+diff(u['i1'],axis=0)/2
# vi1p=v['i1'][:,:-1]+diff(v['i1'],axis=1)/2
# um2p=u['m2'][:-1,:]+diff(u['m2'],axis=0)/2
# vm2p=v['m2'][:,:-1]+diff(v['m2'],axis=1)/2
# ui2p=u['i2'][:-1,:]+diff(u['i2'],axis=0)/2
# vi2p=v['i2'][:,:-1]+diff(v['i2'],axis=1)/2
umag={}
uang={}
umag['m1']=np.sqrt(u['m1']**2+v['m1']**2)*U
uang['m1']=zeros((n,m))
for i in range(n-1):
for j in range(m-1):
uang['m1'][i,j]=math.atan2(v['m1'][i,j],u['m1'][i,j])
umag['i1']=np.sqrt(u['i1']**2+v['i1']**2)*U
uang['i1']=zeros((n,m))
for i in range(n-1):
for j in range(m-1):
uang['i1'][i,j]=math.atan2(v['i1'][i,j],u['i1'][i,j])
umag['m2']=np.sqrt(u['m2']**2+v['m2']**2)*U
uang['m2']=zeros((n,m))
for i in range(n-1):
for j in range(m-1):
uang['m2'][i,j]=math.atan2(v['m2'][i,j],u['m2'][i,j])
umag['i2']=np.sqrt(u['i2']**2+v['i2']**2)*U
uang['i2']=zeros((n,m))
for i in range(n-1):
for j in range(m-1):
uang['i2'][i,j]=math.atan2(v['i2'][i,j],u['i2'][i,j])
upf(umag,uang,'uall',dirname,xmmat,ymmat,30,90,mspo,0)
upf(umag,uang,'unorth',dirname,xmmat,ymmat,6,90,mspo,400)
def ppvt(dirname):
import pickle
from pylab import *
from operators import *
from pf2d_ztitvar import pf2d_ztitvar
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
predir='/data/ilebras/twolayer-jets-climate/'
[psi,q,qfull,u,v,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield_tm.p', 'r' ))
dx=(xmax-xmin)/(n-1)
dy=(ymax-ymin)/(m-1)
dx2=dx**2
tdx=1.0/(2.0*dx)
tdy= 1.0/(2.0*dy)
t4dx=1.0/(8.0*dx)
t4dy=1.0/(8.0*dy)
t4dx2=1.0/(4.0*dx2)
t4dxdy=1.0/(4.0*dx*dy)
t4dy2=1.0/(4.0*dy**2)
twdx2=1.0/(12.0*dx2)
######### calculate ubar grad q #########
dqx1=gradx(qfull['m1'],tdx,n,m)
dqy1=grady(qfull['m1'],tdy,n,m)
dqx2=gradx(qfull['m2'],tdx,n,m)
dqy2=grady(qfull['m2'],tdy,n,m)
pvt={}
pvt['m1']=u['m1']*dqx1+v['m1']*dqy1
pvt['m2']=u['m2']*dqx2+v['m2']*dqy2
######### calculate A del^2 zeta #########
d2xz1=gradxx(q['m1'],t4dx2,n,m)
d2yz1=gradyy(q['m1'],t4dy2,n,m)
d2xz2=gradxx(q['m2'],t4dx2,n,m)
d2yz2=gradyy(q['m2'],t4dy2,n,m)
pvt['i1']=akap*(d2xz1+d2yz1)
pvt['i2']=akap*(d2xz2+d2yz2)
xlow=250
xhigh=-250
ylow=250
pf2d_ztitvar(pvt,'meanpvterms',x,y,dirname,(18,6),xlow,xhigh,ylow,'mean advection term','viscous term',1)
def tracompv(finaldir):
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
whichvar='g'
from readjet2h import readjet2h
vardict=readjet2h(finaldir)
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
f0_nd=f0_dim*L/U
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
[ymat,xmat]=meshgrid(y,x)
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
pload=fortranfile.FortranFile('../testing/'+finaldir+'/hbt.dat')
pvec=pload.readReals()
hb=pvec.reshape(n,m,order='F')
if whichvar=='s':
secspace=0.02/dt
tmaxi=ksnap
elif whichvar=='g':
secspace=0.01
tmaxi=ktmax
time=arange(ksec,tmaxi,1/secspace)*L/U/60**2/24/365.25
tdim=len(time)
#print tdim
xcom2=zeros(tdim)
ycom2=zeros(tdim)
if whichvar=='s':
whiage='a2s'
whiq='q2s'
whip2='p2s'
whip1='p1s'
elif whichvar=='g':
whiage='ag2'
whiq='q2i'
whip2='p2i'
whip1='p1i'
qcom=zeros(tdim)
for iter in range(tdim):
#load some age tracers
pload=fortranfile.FortranFile('../testing/'+finaldir+'/psidat/'+str(whiage)+'_'+str(int(iter)).zfill(3))
pvec=pload.readReals()
age=pvec.reshape(n,m,order='F')
qload=fortranfile.FortranFile('../testing/'+finaldir+'/psidat/'+str(whiq)+'_'+str(int(iter)).zfill(3))
qvec=qload.readReals()
q=qvec.reshape(n,m,order='F')
qload=fortranfile.FortranFile('../testing/'+finaldir+'/psidat/'+str(whip2)+'_'+str(int(iter)).zfill(3))
qvec=qload.readReals()
p2=qvec.reshape(n,m,order='F')
qload=fortranfile.FortranFile('../testing/'+finaldir+'/psidat/'+str(whip1)+'_'+str(int(iter)).zfill(3))
qvec=qload.readReals()
p1=qvec.reshape(n,m,order='F')
qfull=q+fr2*(p1-p2)+f0_nd+betay+hb
qcom[iter]=sum(qfull*age)/sum(age)
figure()
plot(time,qcom,'x-')
title('Tracer center of mass PV')
savefig('../plots/'+finaldir+'_tracerCOMpv'+str(whichvar)+'pos.png')
def qplotall(dirname):
import pickle
from pfunc2d import pfunc2d,plotsame
from pfunc2d_zoom import pfunc2d_zoom
from pylab import *
ioff()
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
predir='/data/ilebras/twolayer-jets-climate/'
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
ymid=y[:-1]+diff(y)/2
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
U=beta_dim*L**2/beta1
# figure(figsize=(13,8))
# contourf(x/1000,y/1000,hbp.T)
# title('Bottom Topography (in m)')
# colorbar()
# savefig(predir+'plots/'+dirname+dirname[6:]+'_hb.png')
fsizemat=(20,6)
xlow=250
xhigh=-250
ylow=250
ytopo=int(m-(ymax-movedw-hw/20)*m/2/ymax-hw)
ytopo2=int(m-(ymax-movedw-hw/20)*m/2/ymax)
#pfunc2d(qfull,'qfull',x,y,dirname,(12,6),mspo,0,optyt1='Mean PV \n y position (km)',optyt2='Instantaneous PV \n y position (km)')
#pfunc2d(psi,'psi_satcomp',x,y,dirname,(12,6),mspo,1,optyt1='Mean $\psi$ \n y position (km)',optyt2='Instantaneous $\psi$ \n y position (km)')
# figure(figsize=(8,5))
# # subplot(121)
# # plotsame(v['m1'][:,m/4:]*U,mspo,x,y[m/4:],1,m/4)
# # plot(x/1e3,y[ytopo]*ones(n)/1e3,'m')
# # plot(x/1e3,y[ytopo2]*ones(n)/1e3,'m')
# # title('Upper Layer')
# # xlabel('x position (km)')
# # ylabel('y position (km)')
# ax=subplot(111)
# plotsame(v['m2'][:,m/3:]*U,mspo,x,y[m/3:],1,m/3)
# plot(x/1e3,y[ytopo]*ones(n)/1e3,'m')
# plot(x/1e3,y[ytopo2]*ones(n)/1e3,'m')
# ax.set_yticklabels([])
# xlabel('x position (km)')
# title('$\overline{v}$')
# colorbar()
# savefig(predir+'plots/'+dirname+dirname[6:]+'_vm2.png',bbox_inches='tight')
#mean lower layer zeta for gsonly
# figure(figsize=(8,5))
# plotsame(q['m2'][:,m/4:]*L/U,mspo,x,y[m/4:],1,m/4)
# xlabel('x position (km)')
# ylabel('y position (km)')
# show()
# savefig(predir+'plots/'+dirname+dirname[6:]+'_zm2.png',bbox_inches='tight')
#mean lower layer zeta for gsslope
# figure(figsize=(10,5))
# ax=subplot(111)
# plotsame(q['m2'][:,m/4:]*L/U,mspo,x,y[m/4:],1,m/4)
# plot(x/1e3,y[ytopo]*ones(n)/1e3,'m')
# plot(x/1e3,y[ytopo2]*ones(n)/1e3,'m')
# xlabel('x position (km)')
# ax.set_yticklabels([])
# show()
# savefig(predir+'plots/'+dirname+dirname[6:]+'_zm2.png',bbox_inches='tight')
import matplotlib
matplotlib.rcParams.update({'font.size':22})
xqind=692
ls=6
figure()
#subplot(211)
contour(x/1e3,y[mspo:]/1e3,psi['m1'][:,mspo:].T,levels=arange(-4e4,4e4+1e4,1e4),colors='k')
# plot(x[1]*ones(m-mspo)/1e3,y[mspo:]/1e3,linewidth=ls+5,color='#a1dab4')
# plot(x[xqind/2]*ones(m-mspo)/1e3,y[mspo:]/1e3,linewidth=ls,color='#41b6c4')
# plot(x[xqind]*ones(m-mspo)/1e3,y[mspo:]/1e3,linewidth=ls,color='#2c7fb8')
#plot(y[mspo:]/1e3,x[xqind*3/4]*ones(m-mspo)/1e3,linewidth=ls,color='#253494')
xlabel('x position (km)')
ylabel('y position (km)')
# subplot(212)
# contour(x,y,psi['m2'].T,levels=arange(-1.5e4,1.5e4+0.5e4,0.5e4),colors='k')
savefig(predir+'plots/'+dirname+dirname[6:]+'_psicartoon.png',bbox_inches='tight')
def loadmeans(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
import os
from operators import *
from uplotfunc import upf
import math
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)
y=arange(ymin,ymax,(ymax-ymin)/m)
dx=(xmax-xmin)/(n-1)
dy=(ymax-ymin)/(m-1)
dx2=dx**2
tdx=1.0/(2.0*dx)
tdy= 1.0/(2.0*dy)
t4dx=1.0/(8.0*dx)
t4dy=1.0/(8.0*dy)
t4dx2=1.0/(4.0*dx2)
t4dxdy=1.0/(4.0*dx*dy)
t4dy2=1.0/(4.0*dy**2)
twdx2=1.0/(12.0*dx2)
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
pload=fortranfile.FortranFile(rundir+'pm2.dat')
pvec=pload.readReals()
pm2=pvec.reshape(n,m,order='F')
psi['m2']=pm2
# #--------and then calculate q (zeta)--------------
q={}
q['m1']=gradxx(psi['m1'],t4dx2,n,m)+gradyy(psi['m1'],t4dy2,n,m)
q['m2']=gradxx(psi['m2'],t4dx2,n,m)+gradyy(psi['m2'],t4dy2,n,m)
#--------if there is bottom topography, plot it-----
try:
pload=fortranfile.FortranFile(rundir+'hbt.dat')
pvec=pload.readReals()
hb=pvec.reshape(n,m,order='F')
except:
hb=zeros((n,m))
hbp=hb
#--------and then plot qfull= q +fr(psi-psi, depending) + betay + hb--------
qfull={}
qfull['m1']=q['m1']*L/U+betay+fr1*(psi['m2']-psi['m1'])+f0_nd
qfull['m2']=q['m2']*L/U+betay+fr2*(psi['m1']-psi['m2'])+hb+f0_nd
#--------calculate a few more things leading up to transport--------
u={}
v={}
u['m1']=-grady(psi['m1'],t4dy,n,m)
u['m2']=-grady(psi['m2'],t4dy,n,m)
v['m1']=gradx(psi['m1'],t4dx,n,m)
v['m2']=gradx(psi['m2'],t4dx,n,m)
umag={}
uang={}
umag['m1']=np.sqrt(u['m1']**2+v['m1']**2)
uang['m1']=zeros((n,m))
for i in range(n-1):
for j in range(m-1):
uang['m1'][i,j]=math.atan2(v['m1'][i,j],u['m1'][i,j])
umag['m2']=np.sqrt(u['m2']**2+v['m2']**2)
uang['m2']=zeros((n,m))
for i in range(n-1):
for j in range(m-1):
uang['m2'][i,j]=math.atan2(v['m2'][i,j],u['m2'][i,j])
umag['i1']=umag['m1']
umag['i2']=umag['m2']
uang['i1']=uang['m1']
uang['i2']=uang['m2']
#xmid=x[:-1]+diff(x)/2
#ymid=y[:-1]+diff(y)/2
[xmmat,ymmat]=meshgrid(x,y)
if not os.path.exists(predir+'pickles/quickfields/'+dirname[:6]):
os.mkdir(predir+'pickles/quickfields/'+dirname[:6])
pickle.dump([psi,q,qfull,u,v,x,y], open(predir+'pickles/quickfields/'+dirname+'_qfield_tm.p', 'w' ) )
upf(umag,uang,'umeantest',dirname,xmmat,ymmat,6,90,mspo,400)
def mtmoc(dirname):
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/"
whivar = "a2s_"
iors = "s"
savedir = "bupu/"
if not os.path.exists("../plots/" + dirname + "_" + whivar + savedir):
os.mkdir("../plots/" + dirname + "_" + whivar + savedir)
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
secspace = 0.1 / dt
time = arange(ksec, ksnap, 1 / secspace) * L / U / 60 ** 2 / 24 / 365.25
# maxrep=len(time)
maxrep = 400
ytopo = int(m - (ymax - movedw - hw / 20) * m / 2 / ymax - hw)
ydw = int(m - (ymax - movedw) * m / 2 / ymax)
ybtwn = (ytopo + ydw) / 2
figure(figsize=(20, 6))
plt.tight_layout
x1 = n / 2 - 500
x2 = n / 2 + 250
fs = 18
tracsum = zeros(maxrep)
tracpos = zeros(maxrep)
tracpos2 = zeros(maxrep)
for iter in range(maxrep):
print iter
aload = fortranfile.FortranFile(predir + pathint + whivar + str(iter).zfill(3))
avec = aload.readReals()
age = avec.reshape(n, m, order="F")
tracsum[iter] = sum((age[x1:x2, :ybtwn] - ybtwn) > 0)
## tracpos[iter]=y[(argmin(abs(age[x1:x2,:].T-ydw),axis=0)).min()]/1e3
## tracpos2[iter]=y[(argmin(abs(age[x1:x2,:].T-ytopo),axis=0)).min()]/1e3
# clf()
# cla()
# contour(x[mspo/2:-mspo]/1e3,y[mspo:]/1e3,age[mspo/2:-mspo,mspo:].T,levels=linspace(ytopo,ydw,4),vmin=ytopo,vmax=ydw,colors='w',linewidth=3)
# age[age>ytopo]=ytopo
# contourf(x[mspo/2:-mspo]/1e3,y[mspo:]/1e3,age[mspo/2:-mspo,mspo:].T,501,cmap='BuPu')
# plot(x[x1]*ones(m)/1e3,y/1e3,'k')
# plot(x[x2]*ones(m)/1e3,y/1e3,'k')
# clim([0, age.max()])
# xlim([x[mspo/2]/1e3, x[-mspo]/1e3])
# ylim([y[mspo]/1e3, y.max()/1e3])
# xlabel('x position (km)',fontsize=fs)
# ylabel('y position (km)',fontsize=fs)
# xlabel(dirname)
# title('Lower Layer')
# savefig(predir+'plots/'+dirname+'_'+whivar+savedir+str(dirname[7:])+'_'+str(int(iter)).zfill(3)+'.png')
# figure()
# plot(time[:maxrep],tracsum)
# #plot(time[:maxrep],tracpos2)
# grid('on')
# savefig(predir+'plots/'+dirname+dirname[6:]+'_tracsumpassed_ybtwn_longer.png',bbox_inches='tight')
tracflux = tracsum.max()
tfluxmn = mean(tracsum)
pickle.dump([tracflux, tfluxmn], open(predir + "pickles/scatter/" + dirname + "_tracvars.p", "w"))
def psimeas(dirname):
import pickle
from pfunc2d import pfunc2d,plotsame
from pfunc2d_zoom import pfunc2d_zoom
ion()
from readjet2h import readjet2h
vardict=readjet2h(dirname)
#print vardict
globals().update(vardict)
predir='/data/ilebras/twolayer-jets-climate/'
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
ymid=y[:-1]+diff(y)/2
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
#beta1=vardict['beta1']
U=beta_dim*L**2/beta1
fsizemat=(20,6)
xlow=250
xhigh=-250
ylow=250
ytopo=int(m-(ymax-movedw-hw/20)*m/2/ymax-hw)
ytopo2=int(m-(ymax-movedw-hw/20)*m/2/ymax)
import matplotlib
matplotlib.rcParams.update({'font.size':22})
xqind=692
ls=6
#find the psi contours that intersect the western and/or eastern boundaries, plot them in a different color than the ones that dont
psicentral1=psi['m1'][(psi['m1']>max(abs(psi['m1'][-1,:]))) & (psi['m1']>max(abs(psi['m1'][:,mspo])))]
#& (psi['m1']>max(psi['m1'][:,-1]))
closedp1=psicentral1.min()
#print closedp
psicentral2=psi['m2'][(psi['m2']>max(abs(psi['m2'][-1,:]))) & (psi['m2']>max(abs(psi['m2'][:,mspo]))) & (psi['m2']>max(psi['m2'][:,-1]))]
closedp2=psicentral2.min()
psibaro=psi['m1']/fr1+psi['m2']/fr2
psicentralbaro=psibaro[(psibaro>max(abs(psibaro[-1,:]))) & (psibaro>max(abs(psibaro[:,mspo])))& (psibaro>max(psibaro[:,-1]))]
closedpb=psicentralbaro.min()
#print closedpb
def conthigh(psiv,closedp,var):
figure(figsize=fsizemat)
contour(x/1e3,y[mspo:]/1e3,psiv[:,mspo:].T,colors='r')
contour(x/1e3,y[mspo:]/1e3,psiv[:,mspo:].T,levels=psiv[0,:],colors='k')
cs=contour(x/1e3,y[mspo:]/1e3,psiv[:,mspo:].T,colors='g',levels=[-closedp,closedp],linewidth=ls)
paths=cs.collections[0].get_paths()
#print paths
plot(x[mspo/2]*ones(m-mspo)/1e3,y[mspo:]/1e3,'k')
plot(x[-mspo]*ones(m-mspo)/1e3,y[mspo:]/1e3,'k')
xlabel('x position (km)')
ylabel('y position (km)')
title(var+', highlighting contours that do not intersect the western boundary')
savefig(predir+'plots/'+dirname+dirname[6:]+'_'+var+'_contour.png',bbox_inches='tight')
return paths
paths1=conthigh(psi['m1'],closedp1,'psi1')
paths2=conthigh(psi['m2'],closedp2,'psi2')
#pathsb=conthigh(psibaro,closedpb,'psibaro')
def maxrecirconly(psiv,closedp,var):
figure(figsize=fsizemat)
cs=contour(x/1e3,y[mspo:]/1e3,psiv[:,mspo:].T,colors='k',levels=[-closedp,closedp],linewidth=ls)
clabel(cs)
xlabel('x position (km)')
ylabel('y position (km)')
title('Largest recirculating streamfunction in '+var)
savefig(predir+'plots/'+dirname+dirname[6:]+var+'_maxrecirc.png',bbox_inches='tight')
# maxrecirconly(psi['m1'],closedp1,'psi1')
# maxrecirconly(psi['m2'],closedp2,'psi2')
# maxrecirconly(psibaro,closedpb,'psibaro')
def linplot(psi,var):
figure(figsize=fsizemat)
plot(y[mspo:]/1e3,psi[0,mspo:].T,'r',label='All streamfunctions')
plot(y[mspo:]/1e3,psi[:,mspo:].T,'r')
plot(y[mspo:]/1e3,psi[-1,mspo:],'k',label='streamfunction on eastern boundary')
plot(y[mspo:]/1e3,psi[0,mspo:],'b',label='streamfunction on western boundary')
xlim([y[mspo]/1e3,y[-1]/1e3])
xlabel('y position (km)')
ylabel(str(var))
grid('on')
legend(loc=(1.05,0.5))
savefig(predir+'plots/'+dirname+dirname[6:]+var+'_psimeas_line.png',bbox_inches='tight')
linplot(psi['m1'],'psi1')
linplot(psi['m2'],'psi2')
linplot(psibaro,'psibaro')
pickle.dump([closedp1,paths1,closedp2,paths2], open( predir+'pickles/psimeas/'+dirname+'_psimeas' , 'w'))
def stabcrit(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
ion()
grid('on')
predir='/data/ilebras/twolayer-jets-climate/'
rundir=predir+'testing/'+dirname+'/'
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
[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
#get beta*y vector to add to
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
where_west=mspo/2+5
where_center=n/2
where_east=3*n/4
from operators import grady
dy=(ymax-ymin)/real(m-1)
t4dy=1.0/(8.0*dy)
dqdy1=grady(qfull['m1'],t4dy,n,m)
dqdy2=grady(qfull['m2'],t4dy,n,m)
ymid=(y[:-1]+y[1:])/2
midy=m/2
figure(figsize=(18,10))
ax=subplot(221)
plot(y[midy:]/1e6,psi['m1'][where_west,midy:],'r')
plot(y[midy:]/1e6,psi['m2'][where_west,midy:],'b')
plot(y[midy:]/1e6,zeros(m-midy),'k')
plot(y[midy:]/1e6,zeros(m-midy),'k.-')
plot(y[midy:]/1e6,zeros(m-midy),'k--')
plot(y[midy:]/1e6,psi['m2'][where_east,midy:],'b--')
plot(y[midy:]/1e6,psi['m1'][where_east,midy:],'r--')
plot(y[midy:]/1e6,psi['m2'][where_center,midy:],'b-.')
plot(y[midy:]/1e6,psi['m1'][where_center,midy:],'r-.')
legend(['Upper Layer','Lower Layer','West (just outside sponge)','Center of domain','East (3/4 of domain)'])
xlim([y[midy]/1e6, y[-1]/1e6])
ax.set_xticklabels([])
title('Psi')
grid('on')
ax=subplot(222)
plot(y[midy:]/1e6,q['m2'][where_west,midy:],'b')
plot(y[midy:]/1e6,q['m1'][where_west,midy:],'r')
plot(y[midy:]/1e6,q['m2'][where_east,midy:],'b--')
plot(y[midy:]/1e6,q['m1'][where_east,midy:],'r--')
plot(y[midy:]/1e6,q['m2'][where_center,midy:],'b-.')
plot(y[midy:]/1e6,q['m1'][where_center,midy:],'r-.')
xlim([y[midy]/1e6, y[-1]/1e6])
ax.set_xticklabels([])
title('Relative vorticity')
grid('on')
ax=subplot(223)
plot(y[midy:]/1e6,qfull['m2'][where_west,midy:],'b')
plot(y[midy:]/1e6,qfull['m1'][where_west,midy:],'r')
plot(y[midy:]/1e6,qfull['m2'][where_east,midy:],'b--')
plot(y[midy:]/1e6,qfull['m1'][where_east,midy:],'r--')
plot(y[midy:]/1e6,qfull['m2'][where_center,midy:],'b-.')
plot(y[midy:]/1e6,qfull['m1'][where_center,midy:],'r-.')
xlim([y[midy]/1e6, y[-1]/1e6])
xlabel('y distance (1000 km)')
title('PV')
grid('on')
subplot(224)
plot(y[midy:]/1e6,dqdy2[where_west,midy:],'b')
plot(y[midy:]/1e6,dqdy1[where_west,midy:],'r')
plot(y[midy:]/1e6,dqdy2[where_east,midy:],'b--')
plot(y[midy:]/1e6,dqdy1[where_east,midy:],'r--')
plot(y[midy:]/1e6,dqdy2[where_center,midy:],'b-.')
plot(y[midy:]/1e6,dqdy1[where_center,midy:],'r-.')
#plot(ymid[midy:]/1e6,zeros(shape(ymid[midy:])),'k')
ylim([-10,10])
xlim([ymid[midy]/1e6, ymid[-1]/1e6])
xlabel('y distance (1000 km)')
title('d(PV)/dy')
grid('on')
savefig(predir+'plots/'+dirname+dirname[6:]+'_stabcrit.pdf')
midy=3*m/4
figure(figsize=(18,10))
ax=subplot(221)
plot(y[midy:]/1e6,psi['m1'][where_west,midy:],'r')
plot(y[midy:]/1e6,psi['m2'][where_west,midy:],'b')
plot(y[midy:]/1e6,zeros(m-midy),'k')
plot(y[midy:]/1e6,zeros(m-midy),'k.-')
plot(y[midy:]/1e6,zeros(m-midy),'k--')
plot(y[midy:]/1e6,psi['m2'][where_east,midy:],'b--')
plot(y[midy:]/1e6,psi['m1'][where_east,midy:],'r--')
plot(y[midy:]/1e6,psi['m2'][where_center,midy:],'b-.')
plot(y[midy:]/1e6,psi['m1'][where_center,midy:],'r-.')
legend(['Upper Layer','Lower Layer','West (just outside sponge)','Center of domain','East (3/4 of domain)'])
xlim([y[midy]/1e6, y[-1]/1e6])
ax.set_xticklabels([])
title('Psi')
grid('on')
ax=subplot(222)
plot(y[midy:]/1e6,q['m2'][where_west,midy:],'b')
plot(y[midy:]/1e6,q['m1'][where_west,midy:],'r')
plot(y[midy:]/1e6,q['m2'][where_east,midy:],'b--')
plot(y[midy:]/1e6,q['m1'][where_east,midy:],'r--')
plot(y[midy:]/1e6,q['m2'][where_center,midy:],'b-.')
plot(y[midy:]/1e6,q['m1'][where_center,midy:],'r-.')
xlim([y[midy]/1e6, y[-1]/1e6])
ax.set_xticklabels([])
title('Relative vorticity')
grid('on')
ax=subplot(223)
plot(y[midy:]/1e6,qfull['m2'][where_west,midy:],'b')
plot(y[midy:]/1e6,qfull['m1'][where_west,midy:],'r')
plot(y[midy:]/1e6,qfull['m2'][where_east,midy:],'b--')
plot(y[midy:]/1e6,qfull['m1'][where_east,midy:],'r--')
plot(y[midy:]/1e6,qfull['m2'][where_center,midy:],'b-.')
plot(y[midy:]/1e6,qfull['m1'][where_center,midy:],'r-.')
xlim([y[midy]/1e6, y[-1]/1e6])
xlabel('y distance (1000 km)')
title('PV')
grid('on')
subplot(224)
plot(y[midy:]/1e6,dqdy2[where_west,midy:],'b')
plot(y[midy:]/1e6,dqdy1[where_west,midy:],'r')
plot(y[midy:]/1e6,dqdy2[where_east,midy:],'b--')
plot(y[midy:]/1e6,dqdy1[where_east,midy:],'r--')
plot(y[midy:]/1e6,dqdy2[where_center,midy:],'b-.')
plot(y[midy:]/1e6,dqdy1[where_center,midy:],'r-.')
#plot(ymid[midy:]/1e6,zeros(shape(ymid[midy:])),'k')
#ylim([-10,10])
xlim([ymid[midy]/1e6, ymid[-1]/1e6])
xlabel('y distance (1000 km)')
title('d(PV)/dy')
grid('on')
savefig(predir+'plots/'+dirname+dirname[6:]+'_stabcrit_zoom.pdf')
def movie_tracer_wpv(dirname):
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
import matplotlib.cm as cmapo
ioff()
predir = "/data/ilebras/twolayer-jets-climate/"
pathint = "testing/" + dirname + "/psidat/"
maxrep = 400
whivar = "a2s_"
iors = "s"
if not os.path.exists("../plots/" + dirname + dirname[6:] + "_" + whivar + "mov"):
os.mkdir("../plots/" + dirname + dirname[6:] + "_" + whivar + "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
# pause time when watching psi movie
ptime = 0.0001
pload = fortranfile.FortranFile(predir + pathint + whivar + "000")
pvec = pload.readReals()
age = pvec.reshape(n, m, order="F")
# set the range for the colorbar to be greater than zero and to the max
cmin1 = age.min()
cmax1 = age.max()
# cmax1=0.5
# cmin=0
# cmax=500
setc1 = linspace(cmin1, cmax1, 151)
hload = fortranfile.FortranFile(predir + "testing/150327/base/hbt.dat")
hvec = hload.readReals()
hb = hvec.reshape(n, m, order="F")
figure(figsize=(20, 6))
# contourf(x/1e3,y/1e3,age.T,cmap=cm.hot_r)#-baseage)#, setc1)
# xlim([x.min()/1e3, x.max()/1e3])
# ylim([y.min()/1e3, y.max()/1e3])
# xlabel(dirname)
# colorbar()
plt.tight_layout
# pause(ptime)
for iter in range(maxrep):
print iter
clf()
aload = fortranfile.FortranFile(predir + pathint + whivar + str(iter).zfill(3))
avec = aload.readReals()
age = avec.reshape(n, m, order="F")
q2load = fortranfile.FortranFile(predir + pathint + "q2" + str(iors) + "_" + str(iter).zfill(3))
q2vec = q2load.readReals()
q2 = q2vec.reshape(n, m, order="F")
p1load = fortranfile.FortranFile(predir + pathint + "p1" + str(iors) + "_" + str(iter).zfill(3))
p1vec = p1load.readReals()
p1 = p1vec.reshape(n, m, order="F")
p2load = fortranfile.FortranFile(predir + pathint + "p2" + str(iors) + "_" + str(iter).zfill(3))
p2vec = p2load.readReals()
p2 = p2vec.reshape(n, m, order="F")
qfull = q2 + betay + fr2 * (p1 - p2) + hb + f0_nd
cla()
contour(x / 1e3, y / 1e3, qfull.T, linspace(8, 14, 6), colors="k")
contourf(x / 1e3, y / 1e3, age.T, setc1, cmap=cmapo.hot_r, vmin=0, vmax=age.max())
clim([0, age.max()])
colorbar()
xlim([x.min() / 1e3, x.max() / 1e3])
ylim([y.min() / 1e3, y.max() / 1e3])
xlabel(dirname)
title(str(iter))
# load instantaneous PV to plot on there too
savefig(
predir
+ "plots/"
+ dirname
+ dirname[6:]
+ "_"
+ whivar
+ "mov/"
+ str(dirname[7:])
+ "_wpv2_"
+ str(int(iter)).zfill(3)
+ ".png"
)
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')
def pvinitnd(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'))
U=beta_dim*L**2/beta1
gprime=f0_dim**2*L**2/h1/fr1
h2=f0_dim**2*L**2/gprime/fr2
#loading/running it is optional, but will make sure I am using the relevant pickle
# from dimensions import dimensions
# dimensions()
# #load the pickle created in dimensions
# with open('/data/ilebras/twolayer-jets-climate/pickles/dim/dimpickle','r') as f:
# [L,xmin,xmax,ymin,ymax,n,m,U,beta_nd,beta_dim,N,fr1,fr2,f0_dim,gprime,h1,h2,gs1,g1,dw1,d1,as2,a2,movedw,movegs,nspo,mspo,hw,hbmax]=pickle.load(f)
#are jsi and jso different? jsdiff=0 if they are not, 1 if they are.
jsdiff=0
x=linspace(xmin,xmax,n)
y=linspace(ymin,ymax,m)
jsi2=zeros((m,1))
jsi1=zeros((m,1))
jso2=zeros((m,1))
jso1=zeros((m,1))
PSI1=zeros((n,m))
PSI2=zeros((n,m))
weig=zeros((n,1))
moveout=0
dwiden=1
for j in range(0,m):
jsi1[j]=gs1*(erf(g1*(y[j]-movegs)))+dw1*(erf(d1*(y[j]-movedw)))
jsi2[j]=as2*(erf(a2*(y[j]-movedw)))
# jso1[j]=gs1*(erf(g1*(y[j]-movegs)))+dw1*(erf(d1*(y[j]-movedw)))
# jso2[j]=as2*(erf(a2/dwiden*(y[j]-movedw-moveout)))
u1=squeeze(diff(jsi1,axis=0))/diff(y)
u2=squeeze(diff(jsi2,axis=0))/diff(y)
# figure()
# plot(y[1:]*L/1e3,u1*U)
# plot(y[1:]*L/1e3,u2*U)
# title('Max GS speed = '+str("%.2f" % u1.min())+'m/s Max DWBC speed = '+str("%.3f" %u1.max())+'m/s')
# grid('on')
if jsdiff==1:
#will need to use this only if jsi is different than jso
for i in range(nspo):
weig[i]= 0.0
for j in range(m):
PSI1[i,j]=jsi1[j]*(1.0-weig[i])+jso1[j]*weig[i]
PSI2[i,j]=jsi2[j]*(1.0-weig[i])+jso2[j]*weig[i]
for i in range(nspo,n-mspo):
weig[i]=float(i-nspo)/(n-(nspo+mspo+1))
for j in range(m):
PSI1[i,j]=jsi1[j]*(1.0-weig[i])+jso1[j]*weig[i]
PSI2[i,j]=jsi2[j]*(1.0-weig[i])+jso2[j]*weig[i]
for i in range(n-mspo,n):
weig[i]=1.0
for j in range(m):
PSI1[i,j]=jsi1[j]*(1.0-weig[i])+jso1[j]*weig[i]
PSI2[i,j]=jsi2[j]*(1.0-weig[i])+jso2[j]*weig[i]
psi1=PSI1
psi2=PSI2
elif jsdiff==0:
psi1=jsi1
psi2=jsi2
# figure()
# subplot(121)
# contourf(PSI1.T)
# colorbar()
# subplot(122)
# contourf(PSI2.T)
# colorbar()
#layer interface height
# hm=f0_dim/gprime*(psi2-psi1)
# # figure()
# # contourf(hm)
# # colorbar()
# hm1=h1-hm
# hm2=h2+hm
#topo
cz=40.0
cm=0.0
hpos=zeros((n,1))
for i in range(n):
hpos[i]=-cm*exp(-(float(i)-float(n)/2)**2/(2*cz**2))+m-(ymax-movedw)*m/2/ymax+hw/2
hb=zeros((n,m))
for i in range(n):
for j in range(int(hw)):
hb[i,int(j+hpos[i]-hw)]=hbmax*float(j)/hw
for j in range(int(hpos[i]),m):
hb[i,j]=hbmax
epsilon=U/(f0_dim*L)
hbdim=hb*h2*epsilon
# figure()
# contourf(hbdim)
# colorbar()
#transport
# mind=-1
# if jsdiff==1:
# tm1=u1*(hm1[mind,:-1]+hm1[mind,1:])/2*L*(xmax-xmin)/n/1e6
# tm2=u2*(hm2[mind,:-1]+hm2[mind,1:])/2*L*(xmax-xmin)/n/1e6
# elif jsdiff==0:
# tm1=u1*(hm1[:-1]+hm1[1:])/2*L*(xmax-xmin)/n/1e6
# tm2=u2*(hm2[:-1]+hm2[1:])/2*L*(xmax-xmin)/n/1e6
# ctm1=cumsum(tm1)
# ctm2=cumsum(tm2)
# figure()
# subplot(221)
# plot(tm1)
# subplot(222)
# plot(tm2)
# subplot(223)
# plot(ctm1)
# subplot(224)
# plot(ctm2)
#PV-PLOT IT IN ALL ITS WAYS
# q_1=del^2*psi_1+beta*y-fr1(psi_1-psi_2)
# q_2=del^2*psi_2+beta*y+fr2(psi_1-psi_2)+h_B
#res=(xmax-xmin)/n*L
ndstep=(xmax-xmin)/n
if jsdiff==1:
d2xp1=diff(psi1,2,0)/ndstep**2
d2yp1=diff(psi1,2,1)/ndstep**2
d2psi1=d2xp1[:,1:-1]+d2yp1[1:-1,:]
d2xp2=diff(psi2,2,0)/ndstep**2
d2yp2=diff(psi2,2,1)/ndstep**2
d2psi2=d2xp2[:,1:-1]+d2yp2[1:-1,:]
slind=[0]
#above-which slice do you want to look at in the plot?
elif jsdiff==0:
d2psi1=squeeze(diff(psi1,2,0)/ndstep**2)
d2psi2=squeeze(diff(psi2,2,0)/ndstep**2)
slind=range(m-2)
xlen=len(x)
ylen=len(y)
#if jsdiff==1:
# betay_dim=f0_dim+beta_dim*tile(y[1:-1],(xlen-2,1))*L
# psidiff=(psi1[1:-1,1:-1]-psi2[1:-1,1:-1])
# q1_dim=d2psi1+betay_dim-fr1*psidiff/L**2
# q2_dim=d2psi2+betay_dim+fr2*psidiff/L**2+f0_dim*hbdim[1:-1,1:-1]/h2
if jsdiff==0:
# betay_dim=f0_dim+beta_dim*y[1:-1]*L
# psidiff=squeeze(psi1[1:-1]-psi2[1:-1])
# q1_dim=tile(d2psi1.T+betay_dim.T-fr1*psidiff.T/L**2,(n-2,1))
# q2_dim=tile(d2psi2.T+betay_dim.T+fr2*psidiff.T/L**2,(n-2,1))+f0_dim*hbdim[1:-1,1:-1]/h2
# #calculate non dimensional qs for comparison/just to check
# q1_nd2=q1_dim*L/U
# q2_nd2=q2_dim*L/U
# q1_nd=tile(d2psi1.T*L/U+betay_dim.T/(U/L)-fr1*psidiff.T/(L*U),(n-2,1))
# q2_nd=tile(d2psi2.T*L/U+betay_dim.T/(U/L)+fr2*psidiff.T/(L*U),(n-2,1))+hb[1:-1,1:-1]
psidiff=squeeze(psi1[1:-1]-psi2[1:-1])
q1_nd3=d2psi1+f0_dim*L/U+beta1*y[1:-1]-fr1*psidiff
q2_nd3=d2psi2+f0_dim*L/U+beta2*y[1:-1]+fr2*psidiff+squeeze(hb[1,1:-1].T)
if not os.path.exists('../pickles/qbases/'+dirname[:6]):
os.mkdir('../pickles/qbases/'+dirname[:6])
#currently, this is for a very specific purpose, not flexible
with open('/data/ilebras/twolayer-jets-climate/pickles/qbases/'+dirname+'_qbase','w') as f:
pickle.dump([d2psi1,d2psi2,psidiff,q1_nd3,q2_nd3],f)
#qs and hb are always 2d, so use a different name for them
# slind_spec=slind[0]
figure(figsize=(20,6))
subplot(121)
plot(d2psi1,'r');
plot(f0_dim*L/U+beta1*y,'b');
plot(-fr1*psidiff,'g');
plot(q1_nd3.T,'k');
xlabel('y position (km)')
title('Layer 1')
subplot(122)
plot(d2psi2,'r');
plot(f0_dim*L/U+beta2*y,'b');
plot(fr2*psidiff,'g');
plot(hb[0,1:-1],'c');
plot(q2_nd3.T,'k');
#tick_params(axis='y',labelleft='off')
xlabel('y position (km)')
title('Layer 2')
legend(('vertical relative vorticity','planetary vorticity','horizontal rel vorticity','topographic contribution','total PV'),loc=2)
savefig('../plots/'+dirname+'/'+dirname[7:]+'_pvinitnd.png')
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
dirlist=['150612/base','150619/gsmeander1','150619/gsmeander2','150619/gscloser1','150619/gscloser2','150625/m2_noslope','150625/m2_slope1','150625/m2_slope2','150625/m2_slope3','150625/m2_slope4','150629/c1_m1','150629/c1_m2','150629/c2_m1','150629/c2_m2','150612/gsonly','150612/dwbconly']
cvec=['k','m','m','g','g','b','b','b','b','r','c','c','c','c','k','k']
mrhines=zeros((2,len(dirlist)))
vqhs=zeros((2,len(dirlist)))
curdist=zeros(len(dirlist))
slopeval=zeros(len(dirlist))
dd=0
for dirname in dirlist:
print dirname
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
[vqh1sum,vqh2sum,xlim,ylow]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/quickfields/'+dirname+'_vqsum.p','r'))
ybnd1=(ymax-movedw)*m/2/ymax-ylow
ybnd2=(ymax-movedw)*m/2/ymax-ylow+hw/2
vqhs[0,dd]=mean(vqh1sum[ybnd1:ybnd2,:])
vqhs[1,dd]=mean(vqh2sum[ybnd1:ybnd2,:])
pvgrad1=mean(qfull['m1'][xlim:-xlim,-(ymax-movedw)*m/2/ymax]-qfull['m1'][xlim:-xlim,-(ymax-movegs)*m/2/ymax])/(y[-(ymax-movedw)*m/2/ymax]-y[-(ymax-movegs)*m/2/ymax])
pvgrad2=mean(qfull['m2'][xlim:-xlim,-(ymax-movedw)*m/2/ymax]-qfull['m2'][xlim:-xlim,-(ymax-movegs)*m/2/ymax])/(y[-(ymax-movedw)*m/2/ymax]-y[-(ymax-movegs)*m/2/ymax])
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 tracom(finaldir):
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
predir='/data/ilebras/twolayer-jets-climate/'
from readjet2h import readjet2h
vardict=readjet2h(finaldir)
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
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
[ymat,xmat]=meshgrid(y,x)
whichvar='s'
if whichvar=='s':
secspace=0.1*dt
tmaxi=ksnap
elif whichvar=='g':
secspace=0.01
tmaxi=ktmax
time=arange(ksec,tmaxi,1/secspace)*L/U/60**2/24/365.25
tdim=len(time)
#print tdim
# xcom1=zeros(tdim)
# ycom1=zeros(tdim)
# for iter in range(tdim):
# #load some age tracers
# pload=fortranfile.FortranFile('../testing/'+finaldir+'/psidat/ag1_'+str(int(iter)).zfill(3))
# pvec=pload.readReals()
# age=pvec.reshape(n,m,order='F')
# xcom1[iter]=sum(age*xmat)/sum(age)
# ycom1[iter]=sum(age*ymat)/sum(age)
xcom2=zeros(tdim)
ycom2=zeros(tdim)
secmom2=zeros(tdim)
if whichvar=='s':
whiage='a2s'
elif whichvar=='g':
whiage='ag2'
tsmall=5
for iter in range(tdim):
#load some age tracers
pload=fortranfile.FortranFile('../testing/'+finaldir+'/psidat/'+str(whiage)+'_'+str(int(iter)).zfill(3))
pvec=pload.readReals()
age=pvec.reshape(n,m,order='F')
xcom2[iter]=sum(age*xmat)/sum(age)
ycom2[iter]=sum(age*ymat)/sum(age)
secmom2[iter]=sum(age*((xmat-xcom2[iter])**2+(ymat-ycom2[iter])**2))/sum(age)
#PLOT THE CENTER OF MASS
figure()
subplot(211)
#plot(time,xcom1/1e3,'bx')
plot(time,xcom2/1e3,'rx-')
ylabel('x position (km)')
title('Tracer center of mass position')
subplot(212)
#plot(time,ycom1/1e3,'bx')
plot(time,ycom2/1e3,'rx-')
ylabel('y position (km)')
xlabel('time (years)')
#legend(['Layer 1','Layer 2'],loc=2)
savefig(predir+'plots/'+finaldir+'_tracerCOM'+str(whichvar)+'pos.png')
#PLOT THE SECOND MOMENT
figure()
plot(time[:tsmall],secmom2[:tsmall]/1e6,'gx-')
title('Tracer second moment')
xlabel('time (years)')
ylabel('$(km ^2)$')
savefig(predir+'plots/'+finaldir+'Tracer_secmom2'+str(whichvar)+'_'+str(tsmall)+'.png')
def qplotall(dirname):
import pickle
from pfunc2d import pfunc2d
from pfunc2d_zoom import pfunc2d_zoom
from pylab import *
ioff()
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
predir='/data/ilebras/twolayer-jets-climate/'
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
ymid=y[:-1]+diff(y)/2
[L,beta_dim,f0_dim,h1]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
# figure(figsize=(13,8))
# contourf(x/1000,y/1000,hbp.T)
# title('Bottom Topography (in m)')
# colorbar()
# savefig(predir+'plots/'+dirname+dirname[6:]+'_hb.png')
fsizemat=(20,6)
xlow=250
xhigh=-250
ylow=250
# pfunc2d(u,'u',x,y,dirname,fsizemat,mspo,1)
# pfunc2d(v,'v',x,y,dirname,fsizemat,mspo,1)
# pfunc2d_zoom(qfull,'qfull',x,y,dirname,fsizemat,xlow,xhigh,ylow,0)
# pfunc2d_zoom(psi,'Psi',x,y,dirname,fsizemat,xlow,xhigh,ylow,1)
# pfunc2d(psi,'Psi',x,y,dirname,fsizemat,mspo,1)
# pfunc2d(q,'zeta',x,y,dirname,fsizemat,mspo,1)
# pfunc2d(qfull,'qfull',x,y,dirname,fsizemat,mspo,0)
# pfunc2d(h,'hinter',x,y,dirname,fsizemat,mspo,0)
#at center
mind=2
tm1=u['m1'][mind,:]*h['m1'][mind,:]*L*(xmax-xmin)/n/1e6
tm2=u['m2'][mind,:]*h['m2'][mind,:]*L*(xmax-xmin)/n/1e6
ctm1=cumsum(tm1)
ctm2=cumsum(tm2)
# tm1=u['m1']*h['m1']*L*(xmax-xmin)/n/1e6
# tm2=u['m2']*h['m2']*L*(xmax-xmin)/n/1e6
# ctm1=cumsum(tm1,axis=1)
# ctm2=cumsum(tm2,axis=1)
figure(figsize=fsizemat)
ax=subplot(221)
plot(y/1000,tm1.T)
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.T)
#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.T)
xlabel('y (1000 km)')
grid(True)
ax=subplot(224)
plot(y/1000,ctm2.T)
xlabel('y (1000 km)')
grid(True)
savefig(predir+'plots/'+dirname+dirname[6:]+'_trans_west.png')
recircpos={}
pvdelta={}
pvgrad2={}
gstren={}
invwid={}
p1max={}
p2max={}
p1min={}
p2min={}
#dd=0
for dd in dirlist:
dirname=dd
print dirname
from readjet2h import readjet2h
vardict=readjet2h(thisdir+dirname)
locals().update(vardict)
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+thisdir+dirname+'_qfield.p', 'r' ))
#so far, three options, nothing, which is prerecirc, center and postrecic (notice bad spelling), recircenter is centered around recirc max
[xqind,xwidth,ecsum_yvar,mcsum_yvar,ecbotyvar,ygs,ydw,p1mx,p2mx,p1mi,p2mi]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/scatter/'+thisdir+dirname+'_fluxes_recircenter.p','r'))
ytopo=int(m-(ymax-movedw-hw/20)*m/2/ymax-hw)
#ypos=(ydw+ytopo)/2
recircpos[dd]=xqind
maxmflux[dd]=mcsum_yvar.max()
maxeflux[dd]=ecsum_yvar.min()
maxeflux_vq[dd]=ecbotyvar.min()
def hov_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/'
if not os.path.exists('../plots/'+dirname+'/qconts'+iors+'mov'):
#os.mkdir('../plots/'+dirname)
os.mkdir('../plots/'+dirname+'/qconts'+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=300
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)
hload=fortranfile.FortranFile(predir+pathint[:-7]+'hbt.dat')
hvec=hload.readReals()
hb=hvec.reshape(n,m,order='F')
figure()
yprofs=zeros((m,tdim))
for iter in range(tdim):
print iter
clf()
q1load=fortranfile.FortranFile(predir+pathint+'q1'+iors+'_'+str(iter).zfill(3))
q1vec=q1load.readReals()
q1=q1vec.reshape(n,m,order='F')
q2load=fortranfile.FortranFile(predir+pathint+'q2'+iors+'_'+str(iter).zfill(3))
q2vec=q2load.readReals()
q2=q2vec.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')
q1full=q1+betay+fr1*(p2-p1)+f0_nd
q2full=q2+betay+fr2*(p1-p2)+hb+f0_nd
# plot(q2.T,'r')
# plot(f0_nd+betay.T,'b')
# plot(hb.T, 'c')
# plot(fr2*(p1-p2).T,'g')
# plot(qfull.T,'m')
ls=arange(9.35,9.45,0.1)
contour(x/1e3,y[ylow:]/1e3,q1full[:,ylow:].T,levels=ls,colors='b')
#plot slope limits
plot(x/1e3,ones(n)*y[-(ymax-movedw)*m/2/ymax-hw/2]/1e3,'r')
plot(x/1e3,ones(n)*y[-(ymax-movedw)*m/2/ymax+hw/2]/1e3,'r')
#plot dwbc limits
plot(x/1e3,ones(n)*y[-(ymax-movedw)*m/2/ymax-2/a2*m/2/ymax]/1e3,'g')
plot(x/1e3,ones(n)*y[-(ymax-movedw)*m/2/ymax+2/a2*m/2/ymax]/1e3,'g')
plot(x[mspo/2]/1000*ones(m-ylow),y[ylow:]/1000,'k')
plot(x[-mspo]/1000*ones(m-ylow),y[ylow:]/1000,'k')
xlabel('x position (km)')
ylabel('y position (km)')
savefig(predir+'plots/'+dirname+'/qconts'+iors+'mov/'+str(dirname[7:])+'_q1conts_'+str(int(iter)).zfill(3)+'.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 getnd(dirname,iors='s'):
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
import matplotlib.cm as cmapo
from operators import gradx
ion()
predir='/data/ilebras/twolayer-jets-climate/'
pathint='testing/'+dirname+'/psidat/'
print dirname
if not os.path.exists('../plots/'+dirname+'/qconts'+iors+'mov'):
#os.mkdir('../plots/'+dirname)
os.mkdir('../plots/'+dirname+'/qconts'+iors+'mov')
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
[L,beta_dim,f0_dim,h1mag]=pickle.load(open('/data/ilebras/twolayer-jets-climate/pickles/dimize.p','r'))
gprime=f0_dim**2*L**2/h1mag/fr1
h2mag=f0_dim**2*L**2/gprime/fr2
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.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=250
xlim=400
yslope1=-(ymax-movedw)*m/2/ymax-hw/2
yslope2=-(ymax-movedw)*m/2/ymax+hw/2
ybnd=yslope1
hb=hbp*f0_dim*L/h2mag/U
#get min/max of PV on slope
q2aslope=qfull['m2'][xlim:-xlim,yslope1:yslope2]
q2smin=q2aslope.min()
q2smax=q2aslope.max()
# figure()
# contourf(x[xlim:-xlim]/1e3,y[yslope1:yslope2]/1e3,q2aslope.T)
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)
vq1sum=zeros((m-ylow,tdim))
vq2sum=zeros((m-ylow,tdim))
have=zeros((2,tdim))
for iter in range(1):#tdim
print iter
clf()
q1load=fortranfile.FortranFile(predir+pathint+'q1'+iors+'_'+str(iter).zfill(3))
q1vec=q1load.readReals()
q1=q1vec.reshape(n,m,order='F')
q2load=fortranfile.FortranFile(predir+pathint+'q2'+iors+'_'+str(iter).zfill(3))
q2vec=q2load.readReals()
q2=q2vec.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')
q1full=q1+betay+fr1*(p2-p1)+f0_nd
q2full=q2+betay+fr2*(p1-p2)+hb+f0_nd
#get v
dx=(xmax-xmin)/real(n-1)
t4dx=1.0/(8.0*dx)
v1=gradx(p1,t4dx,n,m)
v1mid=v1[xlim:-xlim,ylow:]
v2=gradx(p2,t4dx,n,m)
v2mid=v2[xlim:-xlim,ylow:]
#interpolate q to v coords
q1mid=q1full[xlim:-xlim,ylow:]
q2mid=q2full[xlim:-xlim,ylow:]
#hm=f0_dim/gprime*(p2-p1)*U*L
figure()
subplot(211)
plot(v2mid.T)
subplot(212)
plot(q2mid.T)
vq1=v1mid*q1mid
vq2=v2mid*q2mid
# figure()
# plot(vq2)
vq1sum[:,iter]=sum(vq1,axis=0)
vq2sum[:,iter]=sum(vq2,axis=0)
# figure(figsize=(18,6))
# subplot(121)
# plot(time,have[0,:])
# xlabel('time (years)')
# ylabel('mean h on bound')
# subplot(122)
# plot(time,have[1,:])
# xlabel('time (years)')
# figure(figsize=(18,6))
# subplot(121)
# plot(time,vq1sum[-ybnd-ylow,:])
# xlabel('time (years)')
# ylabel('PV Flux')
# subplot(122)
# plot(time,vq2sum[-ybnd-ylow,:])
# xlabel('time (years)')
# savefig(predir+'plots/'+dirname+dirname[6:]+'_vqsum_noabsorh.png')
pickle.dump([vq1sum,vq2sum,xlim,ylow], open(predir+'pickles/quickfields/'+dirname+'_vqsumcorr.p', 'w' ) )
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 pvbudplot(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
import os
ioff()
rcParams['image.cmap'] = 'RdBu'
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
[mf1,betav1,mthi1,mpvf1,ef1,ethi1,epvf1,visc1,sponge1,resid1,mf2,betav2,mthi2,mpvf2,ef2,ethi2,epvf2,visc2,sponge2,topo,bdrag,resid2,hb,hby]=pickle.load(open( predir+'pickles/pvbud/'+dirname+'_pvbud' , 'r'))
# --------plot them-----------
fsizemat=(20,6)
fsizemat2=(20,9)
figure()
subplot(211)
contourf(x,y,hb.T)
colorbar()
subplot(212)
contourf(x,y,hby.T)
colorbar()
# figure()
# plot(y/1e3,topo[mspo/2:-mspo,:].T,'r')
# figure()
# plot(y/1e3,topo[:mspo/2,:].T,'k')
# plot(y/1e3,topo[-mspo:,:].T,'k')
# plot(y/1e3,topo[mspo/2:-mspo,:].T,'r')
# figure(figsize=fsizemat)
# subplot(221)
# plot(y/1e3,mean(topo[:mspo/2,:].T,axis=1),'g')
# plot(y/1e3,mean(topo[-mspo:,:].T,axis=1),'r')
# plot(y/1e3,topo[mspo/2:-mspo,:].T,'b')
# plot(y/1e3,mean(topo[mspo/2:-mspo,:].T,axis=1),'k',linewidth=3)
# ylim([-0.002,0.002])
# title('topographic term')
# subplot(222)
# plot(y/1e3,mean(epvf2[:mspo/2,:].T,axis=1),'g',label='mean west sponge')
# plot(y/1e3,mean(epvf2[-mspo:,:].T,axis=1),'r',label='mean east sponge')
# plot(y/1e3,epvf2[mspo/2:-mspo,:].T,'b')
# plot(y/1e3,epvf2[mspo/2,:].T,'b',label='all interior')
# plot(y/1e3,mean(epvf2[mspo/2:-mspo,:].T,axis=1),'k',linewidth=3,label='mean interior')
# ylim([-0.002,0.002])
# legend(loc=(1.05,-0.5))
# title('eddy pv flux divergence')
# subplot(223)
# plot(y/1e3,mean(visc2[:mspo/2,:].T,axis=1),'g')
# plot(y/1e3,mean(visc2[-mspo:,:].T,axis=1),'r')
# plot(y/1e3,visc2[mspo/2:-mspo,:].T,'b')
# plot(y/1e3,mean(visc2[mspo/2:-mspo,:].T,axis=1),'k',linewidth=3)
# title('viscosity')
# xlabel('y distance (km)')
# ylim([-0.002,0.002])
# subplot(224)
# plot(y/1e3,mean(resid2[:mspo/2,:].T,axis=1),'g')
# plot(y/1e3,mean(resid2[-mspo:,:].T,axis=1),'r')
# plot(y/1e3,resid2[mspo/2:-mspo,:].T,'b')
# plot(y/1e3,mean(resid2[mspo/2:-mspo,:].T,axis=1),'k',linewidth=3)
# title('residual')
# xlabel('y distance (km)')
# suptitle('Layer 2 ')
# ylim([-0.002,0.002])
# savefig(predir+'plots/'+dirname+dirname[6:]+'_toposlice.png',bbox_inches='tight')
#import matplotlib.cm as cm
from pfunc2d import pfunc2d_4in,pfunc2d_6in
############# LAYER 1 ##############################
# # ##### mean fluxes
# pfunc2d_4in(mf1,betav1,mthi1,mpvf1,'pvf1_m',x,y,dirname,fsizemat,mspo,1,optit1='Mean relative vorticity flux divergence',optit2='Mean planetary vorticity flux divergence',optit3='Mean thickness flux divergence',optit4='Total mean PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 1')
# # # ##### eddy fluxes
# pfunc2d_4in(ef1,zeros((n,m)),ethi1,epvf1,'pvf1_e',x,y,dirname,fsizemat,mspo,1,optit1='Eddy relative vorticity flux divergence',optit2='No eddy planetary vorticity flux divergence',optit3='Eddy thickness flux divergence',optit4='Total eddy PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 1')
# # ##### all fluxes
# pfunc2d_6in(mpvf1,epvf1,mpvf1+epvf1,visc1,sponge1,resid1,'pvf1_tot',x,y,dirname,fsizemat2,mspo,1,optit1='Mean PV flux',optit2='Eddy PV flux',optit3='Mean plus eddy',optit4='Viscosity',optit5='Sponge drag',optit6='Residual',suptit='Layer 1')
# # ############# LAYER 2 ##############################
# # ##### mean fluxes
# pfunc2d_6in(mf2,betav2,mthi2,topo,bdrag,mpvf2,'pvf2_m',x,y,dirname,fsizemat,mspo,1,optit1='Mean relative vorticity flux divergence',optit2='Mean planetary vorticity flux divergence',optit3='Mean thickness flux divergence',optit6='Total mean PV flux divergence',optit4='Topographic term',optit5='Bottom Drag',optyt1='position (km)',optyt2='position (km)',suptit='Layer 2')
# ##### eddy fluxes
# pfunc2d_4in(ef2,zeros((n,m)),ethi2,epvf2,'pvf2_e',x,y,dirname,fsizemat,mspo,1,optit1='Eddy relative vorticity flux divergence',optit2='No eddy planetary vorticity flux divergence',optit3='Eddy thickness flux divergence',optit4='Total eddy PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 2')
# ##### all fluxes
# pfunc2d_6in(mpvf2,epvf2,visc2,sponge2,bdrag,resid2,'pvf2_tot',x,y,dirname,fsizemat2,mspo,1,optit1='Mean PV flux divergence',optit2='Eddy PV flux divergence',optit3='Viscosity',optit4='Sponge drag',optit5='Bottom Drag',optit6='Residual',suptit='Layer 2')
# ######################################################################################
# ######################################################################################
# ######################################################################################
# ############################### PLOTTING SPONGE AND INTERIOR #########################
# ######################################################################################
# ######################################################################################
# ######################################################################################
# #######################SEPARATELY!! ##################################################
# ######################################################################################
# ######################################################################################
# ######################################################################################
# ######################################################################################
###FIRST PLOT ALL WITH SPONGE DICTATING THE COLORBAR
spo=1
############# LAYER 1 ##############################
# ##### mean fluxes
pfunc2d_4in(mf1,betav1,mthi1,mpvf1,'pvf1_m_spo',x,y,dirname,fsizemat,spo,1,optit1='Mean relative vorticity flux divergence',optit2='Mean planetary vorticity flux divergence',optit3='Mean thickness flux divergence',optit4='Total mean PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 1')
# # ##### eddy fluxes
pfunc2d_4in(ef1,zeros((n,m)),ethi1,epvf1,'pvf1_e_spo',x,y,dirname,fsizemat,spo,1,optit1='Eddy relative vorticity flux divergence',optit2='No eddy planetary vorticity flux divergence',optit3='Eddy thickness flux divergence',optit4='Total eddy PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 1')
# ##### all fluxes
pfunc2d_6in(mpvf1,epvf1,mpvf1+epvf1,visc1,sponge1,resid1,'pvf1_tot_spo',x,y,dirname,fsizemat2,spo,1,optit1='Mean PV flux divergence',optit2='Eddy PV flux divergence',optit3='Mean plus eddy',optit4='Viscosity',optit5='Sponge drag',optit6='Residual',suptit='Layer 1')
# ############# LAYER 2 ##############################
# ##### mean fluxes
pfunc2d_6in(mf2,betav2,mthi2,topo,bdrag,mpvf2,'pvf2_m_spo',x,y,dirname,fsizemat,spo,1,optit1='Mean relative vorticity flux divergence',optit2='Mean planetary vorticity flux divergence',optit3='Mean thickness flux divergence',optit6='Total mean PV flux divergence',optit4='Topographic term',optit5='Bottom Drag',optyt1='position (km)',optyt2='position (km)',suptit='Layer 2')
##### eddy fluxes
pfunc2d_4in(ef2,zeros((n,m)),ethi2,epvf2,'pvf2_e_spo',x,y,dirname,fsizemat,spo,1,optit1='Eddy relative vorticity flux divergence',optit2='No eddy planetary vorticity flux divergence',optit3='Eddy thickness flux divergence',optit4='Total eddy PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 2')
##### all fluxes
pfunc2d_6in(mpvf2,epvf2,visc2,sponge2,bdrag,resid2,'pvf2_tot_spo',x,y,dirname,fsizemat2,spo,1,optit1='Mean PV flux divergence',optit2='Eddy PV flux divergence',optit3='Viscosity',optit4='Sponge drag',optit5='Bottom Drag',optit6='Residual',suptit='Layer 2')
##############################
### NOW PLOT ONLY THE INTERIOR
##############################
############# LAYER 1 ##############################
# ##### mean fluxes
pfunc2d_4in(mf1[mspo/2:-mspo,mspo:],betav1[mspo/2:-mspo,mspo:],mthi1[mspo/2:-mspo,mspo:],mpvf1[mspo/2:-mspo,mspo:],'pvf1_m_nospo',x[mspo/2:-mspo],y[mspo:],dirname,fsizemat,spo,1,optit1='Mean relative vorticity flux divergence',optit2='Mean planetary vorticity flux divergence',optit3='Mean thickness flux divergence',optit4='Total mean PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 1')
# # ##### eddy fluxes
pfunc2d_4in(ef1[mspo/2:-mspo,mspo:],zeros((n,m))[mspo/2:-mspo,mspo:],ethi1[mspo/2:-mspo,mspo:],epvf1[mspo/2:-mspo,mspo:],'pvf1_e_nospo',x[mspo/2:-mspo],y[mspo:],dirname,fsizemat,spo,1,optit1='Eddy relative vorticity flux divergence',optit2='No eddy planetary vorticity flux divergence',optit3='Eddy thickness flux divergence',optit4='Total eddy PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 1')
# ##### all fluxes
pfunc2d_6in(mpvf1[mspo/2:-mspo,mspo:],epvf1[mspo/2:-mspo,mspo:],mpvf1[mspo/2:-mspo,mspo:]+epvf1[mspo/2:-mspo,mspo:],visc1[mspo/2:-mspo,mspo:],sponge1[mspo/2:-mspo,mspo:],resid1[mspo/2:-mspo,mspo:],'pvf1_tot_nospo',x[mspo/2:-mspo],y[mspo:],dirname,fsizemat2,spo,1,optit1='Mean PV flux divergence',optit2='Eddy PV flux divergence',optit3='Mean plus eddy',optit4='Viscosity',optit5='Sponge drag',optit6='Residual',suptit='Layer 1')
# ############# LAYER 2 ##############################
# ##### mean fluxes
pfunc2d_6in(mf2[mspo/2:-mspo,mspo:],betav2[mspo/2:-mspo,mspo:],mthi2[mspo/2:-mspo,mspo:],topo[mspo/2:-mspo,mspo:],bdrag[mspo/2:-mspo,mspo:],mpvf2[mspo/2:-mspo,mspo:],'pvf2_m_nospo',x[mspo/2:-mspo],y[mspo:],dirname,fsizemat,spo,1,optit1='Mean relative vorticity flux divergence',optit2='Mean planetary vorticity flux divergence',optit3='Mean thickness flux divergence',optit6='Total mean PV flux divergence',optit4='Topographic term',optit5='Bottom Drag',optyt1='position (km)',optyt2='position (km)',suptit='Layer 2')
##### eddy fluxes
pfunc2d_4in(ef2[mspo/2:-mspo,mspo:],zeros((n,m))[mspo/2:-mspo,mspo:],ethi2[mspo/2:-mspo,mspo:],epvf2[mspo/2:-mspo,mspo:],'pvf2_e_nospo',x[mspo/2:-mspo],y[mspo:],dirname,fsizemat,spo,1,optit1='Eddy relative vorticity flux divergence',optit2='No eddy planetary vorticity flux divergence',optit3='Eddy thickness flux divergence',optit4='Total eddy PV flux divergence',optyt1='position (km)',optyt2='position (km)',suptit='Layer 2')
##### all fluxes
pfunc2d_6in(mpvf2[mspo/2:-mspo,mspo:],epvf2[mspo/2:-mspo,mspo:],visc2[mspo/2:-mspo,mspo:],sponge2[mspo/2:-mspo,mspo:],bdrag[mspo/2:-mspo,mspo:],resid2[mspo/2:-mspo,mspo:],'pvf2_tot_nospo',x[mspo/2:-mspo],y[mspo:],dirname,fsizemat2,spo,1,optit1='Mean PV flux divergence',optit2='Eddy PV flux divergence',optit3='Viscosity',optit4='Sponge drag',optit5='Bottom Drag',optit6='Residual',suptit='Layer 2')
import numpy as np
import os
from pylab import *
import fortranfile
import pickle
ion()
figure()
for dd in dirlist:
finaldir=rundir+dd
from readjet2h import readjet2h
vardict=readjet2h(finaldir)
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
f0_nd=f0_dim*L/U
x=arange(xmin,xmax,(xmax-xmin)/n)*L
y=arange(ymin,ymax,(ymax-ymin)/m)*L
[ymat,xmat]=meshgrid(y,x)
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
pload=fortranfile.FortranFile('../testing/'+finaldir+'/hbt.dat')
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 pvcomp(dirname):
from pylab import *
import pickle
from pfunc2d import pfunc2d
from pfunc2d_zoom import pfunc2d_zoom
ioff()
plt.rcParams['image.cmap'] = 'RdYlBu'
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
predir='/data/ilebras/twolayer-jets-climate/'
[psi,q,qfull,u,v,h,hbp,x,y]=pickle.load(open(predir+'pickles/quickfields/'+dirname+'_qfield.p', 'r' ))
[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
hb=hbp/h2/U*f0_dim*L
betay_vec=beta1*y/L
betay=tile(betay_vec,(n,1))
ymid=arange(1,m-1)
ymspo=arange(1,m-mspo-1)
# with open('/data/ilebras/twolayer-jets-climate/pickles/qbases/'+dirname+'_qbase','r') as f:
# [d2psi1,d2psi2,psidiff,q1_base,q2_base]=pickle.load(f)
# qfull_anom={}
# qfull_anom['m1']=qfull['m1'][1:-1,1:-1]-q1_base
# qfull_anom['m2']=qfull['m2'][1:-1,1:-1]-q2_base
# qfull_anom['i1']=qfull['i1'][1:-1,1:-1]-q1_base
# qfull_anom['i2']=qfull['i2'][1:-1,1:-1]-q2_base
fsizemat=(20,6)
# pfunc2d(qfull_anom,'qfull_anom',x[1:-1],y[1:-1],dirname,fsizemat,mspo,1)
# xlow=250
# xhigh=-250
# ylow=250
# pfunc2d_zoom(qfull_anom,'qfull_anom',x[1:-1],y[1:-1],dirname,fsizemat,xlow,xhigh,ylow,1)
# #examine the effect of topography on relative vorticity
# figure(figsize=(20,6))
# subplot(121)
# plot(q['m1'][mspo/2:-mspo,mspo:].T*L/U,'r')
# plot(ymspo,d2psi1[mspo:],'k',lw=2)
# plot(fr1*(psi['m2'][mspo/2:-mspo,mspo:].T-psi['m1'][mspo/2:-mspo,mspo:].T)/(U*L),'g')
# plot(ymspo,-fr1*psidiff[mspo:],'k',lw=2)
# plot(hb[mspo/2:-mspo,mspo:].T/hbmax,'k')
# subplot(122)
# plot(q['m2'][mspo/2:-mspo,mspo:].T*L/U,'r')
# vert,=plot(q['m2'][0,mspo:]*L/U,'r',label='vertical relative vorticity')
# plot(ymspo,d2psi2[mspo:],'k',lw=2)
# plot(fr2*(psi['m1'][mspo/2:-mspo,mspo:].T-psi['m2'][mspo/2:-mspo,mspo:].T)/(U*L),'g')
# horiz,=plot(fr2*(psi['m1'][0,mspo:]-psi['m2'][0,mspo:])/(U*L),'g',label='horizontal rel vorticity')
# plot(ymspo,fr2*psidiff[mspo:],'k',lw=2)
# plot(hb[mspo/2:-mspo,mspo:].T/hbmax,'k')
# lgd=legend(loc=(1.05,0))
# savefig(predir+'plots/'+dirname+dirname[6:]+'_relvorttopo_nosp.png',bbox_extra_artists=(lgd,), bbox_inches='tight')
# #compare before and after of components
figure(figsize=(18,10))
subplot(121)
plot(q['m1'].T*L/U,'r')
#plot(arange(1,m-1),d2psi1,'k',lw=2)
plot(betay.T,'b')
plot(fr1*(psi['m2'].T-psi['m1'].T)/(U*L),'g')
#plot(ymid,-fr1*psidiff,'k',lw=2)
plot(qfull['m1'].T,'m')
plot(qfull['m1'].T-q['m1'].T*L/U,'k')
plot(qfull['m1'].T-q['m1'].T*L/U+fr1*(psi['m1'].T-psi['m2'].T)/(U*L),'y')
#plot(ymid,q1_base.T,'k',lw=2)
subplot(122)
plot(q['m2'].T*L/U,'r')
vert,=plot(q['m2'][0,:]*L/U,'r',label='vertical relative vorticity')
#plot(arange(1,m-1),d2psi2,'k',lw=2)
plot(betay.T,'b')
plan,=plot(betay[0,:],'b',label='planetary vorticity')
plot(fr2*(psi['m1'].T-psi['m2'].T)/(U*L),'g')
horiz,=plot(fr2*(psi['m1'][0,:]-psi['m2'][0,:])/(U*L),'g',label='horizontal rel vorticity')
#plot(ymid,fr2*psidiff,'k',lw=2)
plot(qfull['m2'].T,'m')
tot,=plot(qfull['m2'][0,:],'m',label='total PV')
plot(qfull['m2'].T-q['m2'].T*L/U,'k')
tot2,=plot(qfull['m2'][0,:].T-q['m2'][0,:].T*L/U,'k',label='total PV minus vertical relative vorticity')
plot(qfull['m2'].T-q['m2'].T*L/U-fr2*(psi['m1'].T-psi['m2'].T)/(U*L),'y')
tot3,=plot(qfull['m2'][0,:].T-q['m2'][0,:].T*L/U-fr2*(psi['m1'][0,:].T-psi['m2'][0,:].T)/(U*L),'y', label='total PV minus all relative vorticity')
plot(hb.T,'c')
topo,=plot(hb[0,:],'c',label='topographic contribution')
#plot(ymid,q2_base.T,'k',lw=2)
lgd=legend(loc=(1.05,0))
savefig(predir+'plots/'+dirname+dirname[6:]+'_qmbreakdown_comp.png',bbox_extra_artists=(lgd,), bbox_inches='tight')
def upf(umag,uang,ftit,dirname,xmmat,ymmat,xskip,yskip,sponge,ylow):
from pylab import *
ion()
predir='/data/ilebras/twolayer-jets-climate/'
from readjet2h import readjet2h
vardict=readjet2h(dirname)
locals().update(vardict)
plt.rcParams['image.cmap'] = 'hot_r'
umax=umag['i1'].max()
[m,n]=shape(xmmat)
cnum=51
figure(figsize=(21,8))
ax=subplot(221)
contourf(xmmat[0,:]/1e3,ymmat[ylow:,0]/1e3,umag['m1'][:,ylow:].T,cnum)
colorbar()
quiver(xmmat[ylow::xskip,::yskip]/1e3,ymmat[ylow::xskip,::yskip]/1e3,cos(uang['m1'][::yskip,ylow::xskip].T),sin(uang['m1'][::yskip,ylow::xskip].T))
plot(xmmat[0,sponge/2]/1e3*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,-sponge]/1000*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,:]/1000,ymmat[sponge,0]/1000*ones(n),'k')
ax.set_xticklabels([])
ylim([ymmat[ylow,0]/1e3,ymmat[-1,0]/1e3])
title('Upper Layer')
ylabel('Time means')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax-2/a2*m/2/ymax,0]/1e3,'g')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax+2/a2*m/2/ymax,0]/1e3,'g')
ax=subplot(222)
contourf(xmmat[0,:]/1e3,ymmat[ylow:,0]/1e3,umag['m2'][:,ylow:].T,cnum)
colorbar()
quiver(xmmat[ylow::xskip,::yskip]/1e3,ymmat[ylow::xskip,::yskip]/1e3,cos(uang['m2'][::yskip,ylow::xskip].T),sin(uang['m2'][::yskip,ylow::xskip].T))
plot(xmmat[0,sponge/2]/1e3*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,-sponge]/1000*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,:]/1000,ymmat[sponge,0]/1000*ones(n),'k')
ax.set_xticklabels([])
ax.set_yticklabels([])
ylim([ymmat[ylow,0]/1e3,ymmat[-1,0]/1e3])
title('Lower Layer')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax-2/a2*m/2/ymax,0]/1e3,'g')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax+2/a2*m/2/ymax,0]/1e3,'g')
ax=subplot(223)
contourf(xmmat[0,:]/1e3,ymmat[ylow:,0]/1e3,umag['i1'][:,ylow:].T,cnum)
colorbar()
quiver(xmmat[ylow::xskip,::yskip]/1e3,ymmat[ylow::xskip,::yskip]/1e3,cos(uang['i1'][::yskip,ylow::xskip].T),sin(uang['i1'][::yskip,ylow::xskip].T))
plot(xmmat[0,sponge/2]/1e3*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,-sponge]/1000*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,:]/1000,ymmat[sponge,0]/1000*ones(n),'k')
ylim([ymmat[ylow,0]/1e3,ymmat[-1,0]/1e3])
xlabel('x position')
ylabel('Snapshots')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax-2/a2*m/2/ymax,0]/1e3,'g')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax+2/a2*m/2/ymax,0]/1e3,'g')
ax=subplot(224)
contourf(xmmat[0,:]/1e3,ymmat[ylow:,0]/1e3,umag['i2'][:,ylow:].T,cnum)
colorbar()
quiver(xmmat[ylow::xskip,::yskip]/1e3,ymmat[ylow::xskip,::yskip]/1e3,cos(uang['i2'][::yskip,ylow::xskip].T),sin(uang['i2'][::yskip,ylow::xskip].T))
plot(xmmat[0,sponge/2]/1e3*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,-sponge]/1000*ones(m-ylow),ymmat[ylow:,0]/1000,'k')
plot(xmmat[0,:]/1000,ymmat[sponge,0]/1000*ones(n),'k')
ax.set_yticklabels([])
ylim([ymmat[ylow,0]/1e3,ymmat[-1,0]/1e3])
xlabel('x position')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax-2/a2*m/2/ymax,0]/1e3,'g')
plot(xmmat[0,:]/1e3,ones(n)*ymmat[-(ymax-movedw)*m/2/ymax+2/a2*m/2/ymax,0]/1e3,'g')
savefig(predir+'plots/'+dirname+dirname[6:]+'_'+ftit+'.png')
