adding estNsq.py

so_box_biogeo/diags/estNsq.py

@@ -0,0 +1,80 @@
+import numpy as np
+from MITgcmutils import rdmds, densjmd95
+from mitgcmgrid import loadgrid
+import matplotlib.pyplot as plt
+"""
+    Estimating Brunt-Vaisala frequency using T and S.
+    In case when N$^2$ is needed but do not have 'DRHODR' saved, 
+    one can still estimate it using T and S, 
+    and the appropriate equation of state.
+"""
+#
+# Constants
+#
+rhoconst = 1.035e3
+g = 9.81
+#
+# Load grid files
+# (hspython is available at https://github.com/hajsong/hspython)
+#
+grd = loadgrid('../results', varname=['XC','YC','RC','hFacC','DRC','RF'])
+[nz, ny, nx] = grd.hFacC.shape
+#
+# Compute the stratification frequency using DRHODR. 
+# It is defined at the center of the level.
+#
+dRHOdr = rdmds('ocestrat', 9, rec=0);    # DRHODR is in the first record in "ocestrat"
+Nsq = - dRHOdr*g/rhoconst*grd.mskC
+#
+# Now, estimate Nsq using T and S. 
+# When computing "drhodr" at the interface between tracer cells,
+# density at upper and lower cell is computed using the pressure at the interface. 
+# Nsq is defined at the interface.
+#
+T = rdmds('dynDiag', 9, rec=2)    # THETA is in the third record in "dynDiag"
+S = rdmds('dynDiag', 9, rec=3)    # SALT is in the fourth record in "dynDiag"
+Nsq_TS = np.zeros([nz, ny, nx])
+for k in range(1,nz):
+    press = -rhoconst*g*grd.RF[k]/1e4                   # pressure at the interface
+    # rho at the center of the upper level
+    urho = densjmd95(S[k-1, :, :],T[k-1, :, :], press)
+    # rho at the center of the lower level 
+    lrho = densjmd95(S[k, :, :], T[k, :, :], press)
+    drhodr = (urho-lrho)/grd.DRC[k]*grd.mskC[k, :, :]*grd.mskC[k-1, :, :]
+    Nsq_TS[k, :, :] = -drhodr*g/rhoconst
+#
+# Estimating N$^2$ using "RHOAnoma" is not appropriate 
+# because "RHOAnoma" is computed using the pressure at that level 
+#
+rho = rdmds('ocestrat', 9, rec=1) + rhoconst    # RHOAnoma is in the second record
+Nsq_ra = np.zeros([nz, ny, nx])
+for k in range(1,nz):
+    drhodz = (rho[k-1, :, :] - rho[k, :, :])/grd.DRC[k]\
+             *grd.mskC[k, :, :]*grd.mskC[k-1, :, :]
+    Nsq_ra[k, :, :] = -drhodz*g/rhoconst
+#
+# Check N$^2$
+#
+showz = 5
+showx = 5
+scale = 1e5
+
+Y, Z = np.meshgrid(grd.YC[:, 0], grd.RC[:showz])
+
+f, ax = plt.subplots(1, 3, figsize=(16, 4))
+
+im = ax[0].contourf(Y, Z, Nsq[:showz, :, showx]*scale, np.arange(0, 15.1 , 1), cmap='Reds')
+cb = plt.clabel(im,colors='black',fmt='%3.1f')
+ax[0].set_title('N$^2$ from DRHODR [x 10$^5$ s$^{-1}$]', color='black', fontsize=15)
+ax[0].set_xlabel('latitude')
+ax[0].set_ylabel('depth [m]')
+
+im = ax[1].contourf(Y, Z, Nsq_TS[:showz,:,showx]*scale, np.arange(0,15.1,1), cmap='Reds')
+cb = plt.clabel(im,colors='black',fmt='%3.1f')
+ax[1].set_title('N$^2$ from T and S [x 10$^5$ s$^{-1}$]', color='black', fontsize=15)
+ax[1].set_xlabel('latitude')
+
+im = ax[2].contourf(Y, Z, Nsq_ra[:showz,:,showx]*scale, np.arange(0,15.1,1), cmap='Reds')
+cb = plt.clabel(im,colors='black',fmt='%3.1f')
+ax[2].set_title('N$^2$, from RHOAnoma [x 10$^5$ s$^{-1}$]', color='black', fontsize=15)
+ax[2].set_xlabel('latitude')

so_box_biogeo/diags/mitgcmgrid.py

@@ -0,0 +1,56 @@
+from MITgcmutils import rdmds
+import numpy as np
+
+def loadgrid(gridpath, region=None, varname=None):
+    """
+    [Function]
+    grd=loadgrid(gridname, region=None, varname=None):
+
+    [Description]
+    Read grid files and load them into the 'grd' object
+
+    [Inputs]
+    gridpath : A path to the grid files.
+    region   : A list defining the boundary for the grid.
+               [x0,x1,y0,y1] (default is [0,nx,0,ny])
+    varname  : A list of strings for the grid data to load
+
+    [Output]
+    grd      : An object containing grid data.
+
+    """
+
+    if varname is None:
+        varname = ['XC','YC','RAC','DXC','DYC','hFacC','hFacW','hFacS',\
+                   'Depth','RC','RF','DRC','DRF','XG','YG','RAZ','DXG','DYG']
+
+    class grd(object):
+        for iv, vname in enumerate(varname):
+            if region is None:
+                exec('tmpvar = rdmds("'+gridpath+varname[iv]+'")')
+                tmpvar = tmpvar.squeeze()
+                exec(varname[iv]+' = tmpvar')
+            else:
+                if vname is 'RC' or vname is 'RF' or vname is 'DRC' or vname is 'DRF':
+                    exec('tmpvar = rdmds("'+gridpath+varname[iv]+'")')
+                else: 
+                    exec('tmpvar = rdmds("'+gridpath+varname[iv]+\
+                         '",region = '+str(region)+')')
+                tmpvar = tmpvar.squeeze()
+                exec(varname[iv]+' = tmpvar')
+            if vname == 'hFacC':
+                mskC = hFacC.copy()
+                mskC[mskC==0] = np.nan
+                mskC[np.isfinite(mskC)] = 1.
+            if vname == 'hFacW':
+                mskW = hFacW.copy()
+                mskW[mskW==0] = np.nan
+                mskW[np.isfinite(mskW)] = 1.
+            if vname == 'hFacS':
+                mskS = hFacS.copy()
+                mskS[mskS==0] = np.nan
+                mskS[np.isfinite(mskS)] = 1.
+            del tmpvar
+    del grd.iv,grd.vname
+
+    return grd