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