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 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' ) )
