def quickplot(dirname):
import numpy as np
import fortranfile
from pylab import *
import pickle
ioff()
#set up coords for plotting psi/q fields
xmin=-75.0
xmax=75.0
ymin=-25.0
ymax=25.0
xdim=1501.0
ydim=501.0
#put dimensional scaling constants in
L=50000.0
U=0.875
f0=1e-4
N=0.005
fr1=1.0
fr2=0.25
H1=L*f0/N/fr1
H2=L*f0/N/fr2
H=(H1+H2)/2
gprime=N**2*H
#define x and y (m)
x=arange(xmin,xmax,(xmax-xmin)/xdim)*L
y=arange(ymin,ymax,(ymax-ymin)/ydim)*L
#get beta*y vector to add to
beta=0.05
betay_vec=beta*y/L
betay=tile(betay_vec,(xdim,1))
f0_nd=f0*L/U
predir='/data/ilebras/twolayer-jets-climate/'
rundir=predir+'testing/'+dirname+'/'
#--# ------start with loading enstrophy and other metrics--------------
# #try:
# f = open(rundir+'checkin_noh', 'rb')
# lines = f.readlines()
# # except:
# # f = open(rundir+'nohup.out', 'rb')
# # lines = f.readlines()[:-1]
# all_data = []
# for line in lines:
# list_of_data = map(float, line.split())
# all_data.append(list_of_data)
# all_data = np.array(all_data)
# time=all_data[:,0]
# enst=all_data[:,1]
# ene=all_data[:,2]
# psivar=all_data[:,3]
# #time in years
# t=time*L/U/60**2/24/365.25
#--------move on to load psi--------------
psi={}
pload=fortranfile.FortranFile(rundir+'pm1.dat')
pvec=pload.readReals()
pm1=pvec.reshape(xdim,ydim,order='F')
psi['m1']=pm1*U*L
pload=fortranfile.FortranFile(rundir+'pm2.dat')
pvec=pload.readReals()
pm2=pvec.reshape(xdim,ydim,order='F')
psi['m2']=pm2*U*L
pload=fortranfile.FortranFile(rundir+'p1i.dat')
pvec=pload.readReals()
p1i=pvec.reshape(xdim,ydim,order='F')
psi['i1']=p1i*U*L
pload=fortranfile.FortranFile(rundir+'p2i.dat')
pvec=pload.readReals()
p2i=pvec.reshape(xdim,ydim,order='F')
psi['i2']=p2i*U*L
# #--------and then load q (zeta)--------------
q={}
pload=fortranfile.FortranFile(rundir+'qm1.dat')
pvec=pload.readReals()
qm1=pvec.reshape(xdim,ydim,order='F')
q['m1']=qm1*U/L
pload=fortranfile.FortranFile(rundir+'qm2.dat')
pvec=pload.readReals()
qm2=pvec.reshape(xdim,ydim,order='F')
q['m2']=qm2*U/L
pload=fortranfile.FortranFile(rundir+'q1i.dat')
pvec=pload.readReals()
q1i=pvec.reshape(xdim,ydim,order='F')
q['i1']=q1i*U/L
pload=fortranfile.FortranFile(rundir+'q2i.dat')
pvec=pload.readReals()
q2i=pvec.reshape(xdim,ydim,order='F')
q['i2']=q2i*U/L
# #--------and then load age --------------
# pload=fortranfile.FortranFile(rundir+'psidat/ag2_015')
# pvec=pload.readReals()
# age2=pvec.reshape(xdim,ydim,order='F')
# age2p=age2#*L/U/60**2/24/365.25 #age for plotting is in years
#--------if there is bottom topography, plot it-----
try:
pload=fortranfile.FortranFile(rundir+'hbt.dat')
pvec=pload.readReals()
hb=pvec.reshape(xdim,ydim,order='F')
hbp=hb*H2*U/f0/L
figure(figsize=(13,8))
contourf(x/1000,y/1000,hbp.T)
title('Bottom Topography (in m)')
colorbar()
savefig(predir+'plots/'+dirname+'_hb.png')
except:
hb=zeros((xdim,ydim))
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'])/(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
#-------load the corresponding background PV and calc anomaly--------
with open('/data/ilebras/twolayer-jets-climate/pickles/qbases/'+dirname+'_qbase3','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
predir='/data/ilebras/twolayer-jets-climate/'
figure(figsize=(10,5))
subplot(121)
plot(q['m1'].T*L/U,'r')
plot(d2psi1,'k',lw=2)
plot(f0_nd+betay.T,'b')
plot(fr1*(psi['m2'].T-psi['m1'].T)/(U*L),'g')
plot(-fr1*psidiff,'k',lw=2)
plot(qfull['m1'].T,'m')
plot(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(d2psi2,'k',lw=2)
plot(f0_nd+betay.T,'b')
plan,=plot(f0_nd+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(fr2*psidiff,'k',lw=2)
plot(qfull['m2'].T,'m')
tot,=plot(qfull['m2'][0,:],'m',label='total PV')
plot(hb.T,'c')
topo,=plot(hb[0,:],'c',label='topographic contribution')
plot(q2_base.T,'k',lw=2)
lgd=legend(loc=(1.05,0))
savefig(predir+'plots/'+dirname+'_qmbreakdown_comp.png',bbox_extra_artists=(lgd,), bbox_inches='tight')
xlow=250
xhigh=-250
ylow=250
figure(figsize=(10,5))
subplot(121)
plot(q['m1'][xlow:xhigh,:].T*L/U,'r')
plot(d2psi1,'k',lw=2)
plot(f0_nd+betay[xlow:xhigh,:].T,'b')
plot(fr1*(psi['m2'][xlow:xhigh,:].T-psi['m1'][xlow:xhigh,:].T)/(U*L),'g')
plot(-fr1*psidiff,'k',lw=2)
plot(qfull['m1'][xlow:xhigh,:].T,'m')
plot(q1_base.T,'k',lw=2)
xlim([ylow,ydim])
subplot(122)
plot(q['m2'][xlow:xhigh,:].T*L/U,'r')
vert,=plot(q['m2'][xlow,:]*L/U,'r',label='vertical relative vorticity')
plot(d2psi2,'k',lw=2)
plot(f0_nd+betay[xlow:xhigh,:].T,'b')
plan,=plot(f0_nd+betay[xlow,:],'b',label='planetary vorticity')
plot(fr2*(psi['m1'][xlow:xhigh,:].T-psi['m2'][xlow:xhigh,:].T)/(U*L),'g')
horiz,=plot(fr2*(psi['m1'][xlow,:]-psi['m2'][xlow,:])/(U*L),'g',label='horizontal rel vorticity')
plot(fr2*psidiff,'k',lw=2)
plot(qfull['m2'][xlow:xhigh,:].T,'m')
tot,=plot(qfull['m2'][xlow,:],'m',label='total PV')
plot(hb[xlow:xhigh,:].T,'c')
topo,=plot(hb[xlow,:],'c',label='topographic contribution')
plot(q2_base.T,'k',lw=2)
lgd=legend(loc=(1.05,0))
savefig(predir+'plots/'+dirname+'_qmbreakdown_comp_zoom.png',bbox_extra_artists=(lgd,), bbox_inches='tight')
#--------calculate a few more things leading up to transport--------
u={}
v={}
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
xmid=(x[0:-1]+x[1:])/2
ymid=(y[0:-1]+y[1:])/2
#calculate layer depth fluctuations from psi
hm=f0/gprime*(psi['m2']-psi['m1'])
hi=f0/gprime*(psi['i2']-psi['i1'])
h={}
h['m1']=H1-hm
h['m2']=H2+hm-hbp
h['i1']=H1-hi
h['i2']=H2+hi-hbp
#calculate and plot transport across meridional cross-section in the center(psi*h)
#at the easternmost side
mind=-1
tm1=u['m1'][mind,:]*(h['m1'][mind,:-1]+h['m1'][mind,1:])/2*L*(xmax-xmin)/xdim/1e6
tm2=u['m2'][mind,:]*(h['m2'][mind,:-1]+h['m2'][mind,1:])/2*L*(xmax-xmin)/xdim/1e6
ctm1=cumsum(tm1)
ctm2=cumsum(tm2)
#barotropic and baroclinic streamfunctions
#note that labelling is a little funny, done so for ease of plotting purposes
baro={}
#barotropic
baro['m1']=fr2*psi['m1']+fr1*psi['m2']
baro['i1']=fr2*psi['i1']+fr1*psi['i2']
#baroclinic
baro['m2']=psi['m1']-psi['m2']
baro['i2']=psi['i1']-psi['i2']
#############################################################################
#--------save relevant mean fields to pickle--------------------------------
############################################################################
#############################################################################
#--------primary plotting section-------------------------------------------
############################################################################
from pfunc2d import pfunc2d
from pfunc2d_zoom import pfunc2d_zoom
#slightly messy still, would like to change so that all plotting is done in scripts and plotsize/sponge could be specified there, maybe even called from jet2.h, for example, but that is still a ways off.
fsizemat=(20,6)
sponge=200
pfunc2d(qfull_anom,'qfull_anom',x[1:-1],y[1:-1],dirname,fsizemat,sponge,1)
pfunc2d_zoom(qfull_anom,'qfull_anom',x[1:-1],y[1:-1],dirname,fsizemat,xlow,xhigh,ylow,1)
pfunc2d_zoom(qfull,'qfull',x,y,dirname,fsizemat,xlow,xhigh,ylow,0)
