global_ocean.cs32x15/input.thsice/check_mom_budget.py

a3772ef4e3 Mart*0001 #!/usr/bin/env python
                0002 # -*- coding: iso-8859-15 -*-
                0003 ######################## -*- coding: utf-8 -*-
7c24b74e87 Mart*0004 """Usage: python ./check_mom_budget.py
a3772ef4e3 Mart*0005 """
                0006
                0007 import sys, os
                0008 import matplotlib.pyplot as plt
                0009 import numpy as np
                0010 try:
55a18b4900 Mart*0011     from MITgcmutils import rdmds
a3772ef4e3 Mart*0012 except:
                0013     # this hack to make sure that MITgcmutils.rdmds is available assumes
                0014     # that we are somewhere within the MITgcm/verification directory
                0015     import re
                0016     cwdstr = os.getcwd()
                0017     dirpath = cwdstr[:re.search("verification",cwdstr).start()]
                0018     sys.path.append(os.path.join(dirpath,'utils/python/MITgcmutils') )
                0019     from MITgcmutils import rdmds
                0020
                0021 rDir, nit, deltaT   = "../tr_run.thsice/", 36010, 86400
55a18b4900 Mart*0022 #rDir, nit, deltaT   = "../tr_run.viscA4/", 86405,  3600
a3772ef4e3 Mart*0023 namF, namFs = 'momDiag', 'srfDiag'
                0024
7c24b74e87 Mart*0025 class gridLoader:
                0026     """a hack to mimic a matlab struct"""
                0027     def __init__(self, gDir):
                0028         self.xC = rdmds(os.path.join(gDir,'XC'))
                0029         self.yC = rdmds(os.path.join(gDir,'YC'))
                0030         self.xG = rdmds(os.path.join(gDir,'XG'))
                0031         self.yG = rdmds(os.path.join(gDir,'YG'))
a3772ef4e3 Mart*0032
7c24b74e87 Mart*0033         self.dXc=rdmds(os.path.join(gDir,'DXC'))
                0034         self.dYc=rdmds(os.path.join(gDir,'DYC'))
                0035         self.dXg=rdmds(os.path.join(gDir,'DXG'))
                0036         self.dYg=rdmds(os.path.join(gDir,'DYG'))
a3772ef4e3 Mart*0037
7c24b74e87 Mart*0038         self.dRf=np.squeeze(rdmds(os.path.join(gDir,'DRF')))
a3772ef4e3 Mart*0039
7c24b74e87 Mart*0040         self.rAc=rdmds(os.path.join(gDir,'RAC'))
                0041         self.rAw=rdmds(os.path.join(gDir,'RAW'))
                0042         self.rAs=rdmds(os.path.join(gDir,'RAS'))
                0043         self.rAz=rdmds(os.path.join(gDir,'RAZ'))
a3772ef4e3 Mart*0044
7c24b74e87 Mart*0045         self.hFacC=rdmds(os.path.join(gDir,'hFacC'))
                0046         self.hFacW=rdmds(os.path.join(gDir,'hFacW'))
                0047         self.hFacS=rdmds(os.path.join(gDir,'hFacS'))
                0048         self.depth=rdmds(os.path.join(gDir,'Depth'))
a3772ef4e3 Mart*0049
                0050
                0051 def readdiags(fname,nit):
                0052     v,iter,M=rdmds(os.path.join(rDir,fname), nit, returnmeta = True);
                0053     fList=M['fldlist']
                0054     n=len(fList)
                0055     return v, n, fList
                0056
                0057 def split_C_cub(v3d):
55a18b4900 Mart*0058     """ v6t = split_C_cub(v3d)
a3772ef4e3 Mart*0059     --------------------------------------------
                0060     split 2d/3d arrays V, center, to 2d/3d x 6 faces
                0061     and add 1 column + 1 row <== at the begining !!!
                0062     => output is v6t[nr,ny+1,ny+1,6]
                0063     """
                0064     kad=1
                0065     dims = v3d.shape
                0066     nx,ny = dims[-1],dims[-2]
                0067     if len(dims) == 2: nr = 1
                0068     else: nr = np.prod(dims[:-2])
                0069     nyp=ny+1; n2p=ny+2; nye=ny+kad
                0070
                0071     v = v3d.reshape((nr,ny,nx))
                0072     v6t = np.zeros((nr,nye,nye,6))
                0073
                0074     for n in range(6):
                0075         v6t[:,1:,1:,n]=v[:,:,n*ny:(n+1)*ny]
                0076
                0077     v6t[:,1:,0, 0]=v6t[:,-1,:0:-1,4]
                0078     v6t[:,1:,0, 2]=v6t[:,-1,:0:-1,0]
                0079     v6t[:,1:,0, 4]=v6t[:,-1,:0:-1,2]
                0080     v6t[:,1:,0, 1]=v6t[:,1:,-1,0]
                0081     v6t[:,1:,0, 3]=v6t[:,1:,-1,2]
                0082     v6t[:,1:,0, 5]=v6t[:,1:,-1,4]
7c24b74e87 Mart*0083
a3772ef4e3 Mart*0084     v6t[:,0, :, 0]=v6t[:,-1,:,5]
                0085     v6t[:,0, :, 2]=v6t[:,-1,:,1]
                0086     v6t[:,0, :, 4]=v6t[:,-1,:,3]
                0087     v6t[:,0,1:, 1]=v6t[:,:0:-1,-1,5]
                0088     v6t[:,0,1:, 3]=v6t[:,:0:-1,-1,1]
                0089     v6t[:,0,1:, 5]=v6t[:,:0:-1,-1,3]
                0090
                0091     v6t[:,0,0,1]=v6t[:,1,-1,0]
55a18b4900 Mart*0092     v6t[:,0,0,3]=v6t[:,1,-1,2]
                0093     v6t[:,0,0,5]=v6t[:,1,-1,4]
a3772ef4e3 Mart*0094
                0095     return v6t
                0096
                0097 def calc_grad(fld,dx,dy):
55a18b4900 Mart*0098     """calculate gradient of 6-tiled fields fld[nr,ny+1,ny+1,6]
                0099     and return as (dfx[nr,ny,ncx], dfy[nr,ny,ncx])
a3772ef4e3 Mart*0100     """
                0101     nnr, np1 = fld.shape[:2]
                0102     nc = np1-1
                0103     myshape = (1,G.dXc.shape[0],G.dXc.shape[1])
                0104     dfx = (fld[:,:,1:,:]-fld[:,:,:-1,:])[:,1:,:,:] \
                0105           .reshape((nnr,nc,nc*6),order='F') \
                0106           /np.tile(dx.reshape(myshape),(nnr,1,1))
                0107     dfy = (fld[:,1:,:,:]-fld[:,:-1,:,:])[:,:,1:,:] \
                0108           .reshape((nnr,nc,nc*6),order='F') \
                0109           /np.tile(dy.reshape(myshape),(nnr,1,1))
                0110     return dfx, dfy
7c24b74e87 Mart*0111
a3772ef4e3 Mart*0112 def getListIndex(diagName,fldLst):
7c24b74e87 Mart*0113     """Return index of diagName in fldlst (list of diagnostics names);
55a18b4900 Mart*0114     if diagName is not in fldlist, return -1
a3772ef4e3 Mart*0115     """
                0116     if diagName in str(fldLst):
                0117         j = fldLst.index(diagName)
                0118     else:
                0119         j = -1
                0120     return j
                0121
55a18b4900 Mart*0122 def printstats(var,varName):
a3772ef4e3 Mart*0123     fmt='Var = %8s : Min,Max,Avr,StD= %12.5e %12.5e %12.5e %12.5e'
55a18b4900 Mart*0124     print(fmt % (varName, var.min(), var.max(), var.mean(), var.std()))
a3772ef4e3 Mart*0125     return
                0126
                0127 def printsum(var,res):
                0128     fmt = '   Sum Tend: Avr,StD= %12.5e %12.5e ; Residual= %12.5e %12.5e +';
                0129     print(fmt % (var.mean(), var.std(), res.mean(), res.std()))
                0130     return
                0131
                0132 def printStatsAndSum(fldLst,dtot,gtot):
55a18b4900 Mart*0133     """For each of the terms in fldLst
                0134       a) print some stats of this term
                0135       b) add to other tendency and print stats of the sum
                0136       c) substract the sum from total tendency (-> residual) and print stats
                0137     """
a3772ef4e3 Mart*0138     for fldName in fldLst:
                0139         j = getListIndex(fldName,fldList)
7c24b74e87 Mart*0140         if j > -1:
a3772ef4e3 Mart*0141             var = np.copy(np.squeeze(v4d[j,:,:,:]))
55a18b4900 Mart*0142         elif "m_ImplD" in fldName:
a3772ef4e3 Mart*0143             # U/Vm_ImpD was not found. Now we have to do something different
55a18b4900 Mart*0144             print(fldName+" was not found, trying alternative",end = " ")
                0145             if 'Um' in fldName:
a3772ef4e3 Mart*0146                 if juNz>-1: j, var = juNz, gUnuZ
55a18b4900 Mart*0147             elif 'Vm' in fldName:
a3772ef4e3 Mart*0148                 if jvNz>-1: j, var = jvNz, gVnuZ
                0149             if j==-1: print("... unsuccessfull")
                0150
                0151         if j>-1:
                0152             printstats(var,fldList[j])
                0153             gtot=gtot+var
                0154             printsum(gtot,dtot-gtot)
                0155         else:
                0156             print('... cannot use '+fldName)
                0157     return
7c24b74e87 Mart*0158
a3772ef4e3 Mart*0159 # let's go
                0160
55a18b4900 Mart*0161 # first establish the grid parameters that we are going to need
7c24b74e87 Mart*0162 G = gridLoader(rDir)
a3772ef4e3 Mart*0163
                0164 nr, nc = G.hFacC.shape[:2]
                0165 nPxy = G.hFacC.shape[2]*nc
                0166 nPp2 = nPxy+2;
                0167 ncx  = 6*nc
                0168 np1  = nc+1;
                0169
                0170 mskW=np.minimum(np.ceil(G.hFacW),1); mskS=np.minimum(np.ceil(G.hFacS),1);
                0171
55a18b4900 Mart*0172 # set constants
a3772ef4e3 Mart*0173 rhoConst=1035.
                0174 gravity =9.81
                0175
                0176 # Read in 2-D diagnostics file "namFs" and 3-D diagnostics file "namF":
                0177 v3d,nV2d,f2dList = readdiags(namFs,nit)
                0178 v4d,nV,  fldList = readdiags(namF,nit)
                0179
                0180 if nV2d == 0: f2dList=fldList
                0181
55a18b4900 Mart*0182 # compute the horizontal pressure gradient terms in case we need them
a3772ef4e3 Mart*0183 if 'PHI_SURF' in str(f2dList):
                0184     jdps = f2dList.index('PHI_SURF')
                0185     var = v3d[jdps,:,:]
                0186 elif 'ETAN' in str(f2dList):
                0187     jdps = f2dList.index('ETAN')
                0188     var = gravity*np.copy(v3d[jdps,:,:])
                0189 else:
                0190     jdps = -1
7c24b74e87 Mart*0191
55a18b4900 Mart*0192 if jdps > -1:
a3772ef4e3 Mart*0193     dpx, dpy = calc_grad(- split_C_cub(var),G.dXc,G.dYc)
                0194
55a18b4900 Mart*0195 # horizontal non-hydrostatic pressure gradients terms
a3772ef4e3 Mart*0196 jnh = -1
                0197 fileName='%s.%10.10i.%s' % (os.path.join(rDir,'pnhDiag'),nit+1,'data')
                0198 if os.path.isfile(fileName):
                0199     print('  -- loading file: %s ...' % fileName)
                0200     var=rdmds(os.path.join(rDir,'pnhDiag'),nit+1)
                0201     print(' done')
                0202 elif 'PHI_NH' in str(fldList):
                0203     jnh = fldList.index('PHI_NH')
                0204     var=np.copy(v4d[jnh,:,:,:])
                0205
                0206 if jnh > -1:
                0207     dpNHx, dpNHy = calc_grad(- split_C_cub(var),G.dXc,G.dYc)
                0208     dpNHx=dpNHx*mskW
                0209     dpNHy=dpNHy*mskS
                0210
55a18b4900 Mart*0211 # when using z* with  older output, need to account for
                0212 # column vertical streaching in computation of vertical
                0213 # viscosity tendency form vertical viscous flux 'VISrI_Um'
a3772ef4e3 Mart*0214 if 'ETAN' in str(f2dList):
                0215     jeta = f2dList.index('ETAN')
                0216     v6t = split_C_cub(v3d[jeta,:,:]*G.rAc)
                0217     vbx = 0.5*(v6t[:,:,1:,:]+v6t[:,:,:-1,:])[:,1:,:,:] \
                0218            .reshape((1,nc,nc*6),order='F')/G.rAw
                0219     vby = 0.5*(v6t[:,1:,:,:]+v6t[:,:-1,:,:])[:,:,1:,:] \
                0220            .reshape((1,nc,nc*6),order='F')/G.rAs
                0221     d6t = split_C_cub(G.depth)
                0222     hhx = np.minimum(d6t[:,:,1:,:],d6t[:,:,:-1,:])[:,1:,:,:] \
                0223             .reshape((1,nc,nc*6),order='F')
                0224     hhy = np.minimum(d6t[:,1:,:,:],d6t[:,:-1,:,:])[:,:,1:,:] \
                0225             .reshape((1,nc,nc*6),order='F')
55a18b4900 Mart*0226     hhx[hhx==0.]=np.Inf
                0227     rFacW=vbx/hhx + np.ones((nc,ncx))
                0228     hhy[hhy==0.]=np.Inf
                0229     rFacS=vby/hhy + np.ones((nc,ncx))
a3772ef4e3 Mart*0230 else:
                0231     jdps = -1
                0232
                0233
                0234 #-------------------------------------------------------------------------------
                0235
55a18b4900 Mart*0236 # horizontal gradients of the potential Phi, this can be derived from
                0237 # different diagnostics depending on their availability
                0238 gUdp, gVdp = np.zeros((nr,nc,ncx)), np.zeros((nr,nc,ncx))
                0239 titUdp, titVdp = ' ? ', ' ? '
                0240 jdph=-1
a3772ef4e3 Mart*0241 j1=getListIndex('Um_dPhiX',fldList)
                0242 j2=getListIndex('Vm_dPhiY',fldList)
                0243 if j1==-1 & j2==-1:
                0244     jdph=-1;
                0245     j1=getListIndex('Um_dPHdx',fldList)
                0246     j2=getListIndex('Vm_dPHdy',fldList)
                0247     if   j1>-1: jdph=j1
                0248     elif j2>-1: jdph=j2
                0249     if jdph > -1 & jdps > -1:
                0250         gUdp=dpx.repmat((nr,1,1))*mskW;
                0251         gVdp=dpy.repmat((nr,1,1))*mskS;
                0252     if jnh > -1: gUdp=gUdp+dpNHx; gVdp=gVdp+dpNHy
                0253 else:
                0254   if j1==-1: jdph=j2
55a18b4900 Mart*0255   else:      jdph=j1
a3772ef4e3 Mart*0256
                0257 if jdph > -1:
                0258   if j1 > -1:
                0259       gUdp=gUdp+np.squeeze(v4d[j1,:,:,:])
                0260       titUdp=fldList[j1]
                0261   if j2 > -1:
                0262       gVdp=gVdp+np.squeeze(v4d[j2,:,:,:])
                0263       titVdp=fldList[j2]
                0264   if jdps > -1:
                0265     titUdp=titUdp[:-1]+titUdp[-1].upper()
                0266     titVdp=titVdp[:-1]+titVdp[-1].upper()
55a18b4900 Mart*0267     #print(' titUdp: >%s< ; titVdp: >%s<\n' %(titUdp,titVdp))
a3772ef4e3 Mart*0268
                0269 #-- Tendencies from implicit vertical viscous fluxes
                0270 # Note: will be used to close momentum budget
                0271 #   a) if using older output (since 'Um_ImplD' was not there);
                0272 #   b) and using implicit viscosity but without implicit bottom friction
                0273 # In the latest case (selectImplicitDrag=2,) cannot close the budget
                0274 # using older output
                0275 juNz = getListIndex('VISrI_Um',fldList)
55a18b4900 Mart*0276 jvNz = getListIndex('VISrI_Vm',fldList)
                0277 if juNz>-1 or jvNz>-1:
                0278     print('  --  Tendencies from vertically visc. fluxes --')
a3772ef4e3 Mart*0279 if juNz>-1:
                0280     var=np.copy(np.squeeze(v4d[juNz,:,:,:]))
                0281     # compute tendency from minus div of vert. fluxes:
                0282     div=np.copy(var); div[:-1,:,:]=div[:-1,:,:]-var[1:,:,:]
                0283     ddz=G.hFacW*np.tile(G.dRf.reshape((nr,1,1)),(1,nc,ncx))
                0284     rdz=np.copy(ddz); rdz[ddz!=0]=1./rdz[ddz!=0]
                0285     gUnuZ= - div*rdz/np.tile(G.rAw*rFacW,(nr,1,1))
                0286     printstats(gUnuZ,fldList[juNz])
                0287     #--
                0288     jj = getListIndex('Um_ImplD',fldList)
                0289     if jj > -1:
                0290         var=np.copy(np.squeeze(v4d[jj,:,:,:]))
                0291         printstats(var,fldList[jj])
                0292         var = var - gUnuZ
                0293         printstats(var,'Diff:2-1')
                0294
                0295     print()
                0296
                0297 if jvNz>-1:
                0298     var=np.copy(np.squeeze(v4d[jvNz,:,:,:]))
                0299     # compute tendency from minus div of vert. fluxes:
                0300     div=np.copy(var); div[:-1,:,:]=div[:-1,:,:]-var[1:,:,:]
                0301     ddz=G.hFacS*np.tile(G.dRf.reshape((nr,1,1)),(1,nc,ncx))
                0302     rdz=np.copy(ddz); rdz[ddz!=0]=1./rdz[ddz!=0]
                0303     gVnuZ= - div*rdz/np.tile(G.rAs*rFacS,(nr,1,1))
                0304     printstats(gVnuZ,fldList[jvNz])
                0305     #--
                0306     jj = getListIndex('Vm_ImplD',fldList)
55a18b4900 Mart*0307     if jj > -1:
a3772ef4e3 Mart*0308         var=np.copy(np.squeeze(v4d[jj,:,:,:]))
                0309         printstats(var,fldList[jj])
                0310         var = var - gVnuZ
                0311         printstats(var,'Diff:2-1')
                0312
                0313     print()
                0314
                0315 # Here we check that vertical integral of implicit vertical viscous tendency
                0316 # match either bottom drag (if using implicit bottom drag) or simply zero.
                0317 j1 = getListIndex('Um_ImplD',fldList)
                0318 j2 = getListIndex('botTauX' ,f2dList)
                0319 if j1>-1 & j2>-1:
                0320     print('  --  Vertically integrated tendencies --');
                0321     bTauX = np.copy(v3d[j2,:,:])
                0322     printstats(bTauX,f2dList[j2])
                0323     var=np.copy(np.squeeze(v4d[j1,:,:,:]))
                0324     # vertical integration:
                0325     ddz=G.hFacW*np.tile(G.dRf.reshape((nr,1,1)),(1,nc,ncx))
                0326     var=rhoConst*((var*ddz).sum(axis=0))*rFacW
                0327     printstats(var,fldList[j1])
                0328     printstats(var-bTauX,'Diff:2-1')
                0329     print()
                0330
                0331 j1 = getListIndex('Vm_ImplD',fldList)
                0332 j2 = getListIndex('botTauY' ,f2dList)
                0333 if j1>-1 & j2>-1:
                0334     bTauY = np.copy(v3d[j2,:,:])
                0335     printstats(bTauY,f2dList[j2])
                0336     var=np.copy(np.squeeze(v4d[j1,:,:,:]))
                0337     # vertical integration:
                0338     ddz=G.hFacS*np.tile(G.dRf.reshape((nr,1,1)),(1,nc,ncx))
                0339     var=rhoConst*((var*ddz).sum(axis=0))*rFacS
                0340     printstats(var,fldList[j1])
                0341     printstats(var-bTauY,'Diff:2-1')
                0342     print()
                0343
55a18b4900 Mart*0344 # this is where the actual momentum budget check starts
a3772ef4e3 Mart*0345 print('  --  Check Mom budget, exp: %s, files: %s & %s, it= %i'
                0346       % (rDir,namF,namFs,nit))
                0347
                0348 j = getListIndex('TOTUTEND',fldList)
                0349 if j > -1:
                0350     dUtot=np.copy(np.squeeze(v4d[j,:,:,:]))/86400.
                0351     printstats(dUtot,fldList[j])
7c24b74e87 Mart*0352
a3772ef4e3 Mart*0353 j = getListIndex('Um_dPhiX',fldList)
2603069db6 Mart*0354 if j == -1: j = getListIndex('Um_dPHdx',fldList)
a3772ef4e3 Mart*0355 if j >-1:
                0356   printstats(gUdp,titUdp)
                0357   gUtot=gUdp
                0358
55a18b4900 Mart*0359 printStatsAndSum(['Um_Advec','Um_Ext','Um_Diss','Um_ImplD','AB_gU'],dUtot,gUtot)
a3772ef4e3 Mart*0360
                0361 print()
                0362
                0363 j = getListIndex('TOTVTEND',fldList)
                0364 if j > -1:
                0365     dVtot=np.copy(np.squeeze(v4d[j,:,:,:]))/86400.
                0366     printstats(dVtot,fldList[j])
                0367
                0368 j = getListIndex('Vm_dPhiY',fldList)
2603069db6 Mart*0369 if j == -1: j = getListIndex('Vm_dPHdy',fldList)
a3772ef4e3 Mart*0370 if j > -1:
                0371   printstats(gVdp,titVdp)
                0372   gVtot=gVdp
                0373
55a18b4900 Mart*0374 printStatsAndSum(['Vm_Advec','Vm_Ext','Vm_Diss','Vm_ImplD','AB_gV'],dVtot,gVtot)