import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
saveF = False
#-- directories --
dir_in = '/home2/datawork/qjamet/CROCO/RUNS/BASIN'
dir_fig = '/home2/datahome/qjamet/Figures/BASIN'
#- beta-plane -
da1 = xr.open_dataset( \
str("%s/run_%04.i/basin_his.nc" % (dir_in, 5)) )
daM = xr.open_dataset( \
str("%s/run_%04.i/basin_diaM_avg.nc" % (dir_in, 5)) )
daVrt = xr.open_dataset( \
str("%s/run_%04.i/basin_diags_vrt_avg.nc" % (dir_in, 5)) )
#daKE = xr.open_dataset( \
# str("%s/run_%04.i/basin_diags_ek_avg.nc" % (dir_in, 5)) )
#- f-plane -
da2 = xr.open_dataset( \
str("%s/run_%04.i/basin_his.nc" % (dir_in, 6)) )
daM2 = xr.open_dataset( \
str("%s/run_%04.i/basin_diaM_avg.nc" % (dir_in, 6)) )
daVrt2 = xr.open_dataset( \
str("%s/run_%04.i/basin_diags_vrt_avg.nc" % (dir_in, 6)) )
#daKE = xr.open_dataset( \
# str("%s/run_%04.i/basin_diags_ek_avg.nc" % (dir_in, 6)) )
fig1 = plt.figure(figsize=(12, 5))
#
ax1 = fig1.add_subplot(1,2,1)
cs1 = ax1.contourf(da1.x_rho, da1.y_rho, da1.zeta[-1, :, :], 20, cmap='RdBu_r')
plt.colorbar(cs1)
ax1.set_title('beta-plan')
#
ax2 = fig1.add_subplot(1,2,2)
cs2 = ax2.contourf(da2.x_rho, da1.y_rho, da2.zeta[-1, :, :], 20, cmap='RdBu_r')
plt.colorbar(cs2)
ax2.set_title('f-plan')
if saveF:
figN = str("BASIN_zeta_f-plan_beta-plan")
fig1.savefig(str("%s/%s.png" % (dir_fig, figN)), bbox_inches='tight')
fig1.savefig(str("%s/%s.pdf" % (dir_fig, figN)), bbox_inches='tight')
From https://www.jgula.fr/Croco/diagnostics_croco.pdf
vrt_rate = vrt_xadv + vrt_yadv + vrt_prsgrd + vrt_cor + vrt_vmix + vrt_hmix + vrt_nudg
daVrt
ttt = -1
#- beta-plane -
vrt_rate = daVrt.vrt_rate[ttt, :, :]
vrt_rhs = daVrt.vrt_xadv[ttt, :, :] + daVrt.vrt_yadv[ttt, :, :] + daVrt.vrt_prsgrd[ttt, :, :] + \
daVrt.vrt_cor[ttt, :, :] + daVrt.vrt_vmix[ttt, :, :] + daVrt.vrt_hmix[ttt, :, :]
#- f-plane (beta=0) -
vrt2_rate = daVrt2.vrt_rate[ttt, :, :]
vrt2_rhs = daVrt2.vrt_xadv[ttt, :, :] + daVrt2.vrt_yadv[ttt, :, :] + daVrt2.vrt_prsgrd[ttt, :, :] + \
daVrt2.vrt_cor[ttt, :, :] + daVrt2.vrt_vmix[ttt, :, :] + daVrt2.vrt_hmix[ttt, :, :]
jjj = 20
scaleF = 1e3
fig1 = plt.figure(figsize=(12, 10))
fig1.clf()
#
ax1 = fig1.add_subplot(2, 2, 1)
cs1 = ax1.contourf(vrt_rate, 20, cmap='RdBu_r')
plt.colorbar(cs1)
ax1.plot(np.arange(daVrt.dims['xi_u']), jjj*np.ones(daVrt.dims['xi_u']), 'k.')
ax1.set_title(r'1/ $\partial_t (\overline{PV}^{z})$')
#
ax2 = fig1.add_subplot(2, 2, 2)
cs2 = ax2.contourf(vrt_rhs, 20, cmap='RdBu_r')
plt.colorbar(cs2)
ax2.set_title(r'2/ $\sum RHS$')
#
ax3 = fig1.add_subplot(2, 2, 3)
cs3 = ax3.contourf(vrt_rate-vrt_rhs, 20, cmap='RdBu_r')
plt.colorbar(cs3)
ax3.set_title(r'Residual (1/-2/)')
#
ax4 = fig1.add_subplot(2, 2, 4)
p0 = ax4.plot(vrt_rate[jjj, :], 'k', lw=3)
p1 = ax4.plot(vrt_rhs[jjj, :], 'r.')
p2 = ax4.plot((vrt_rate[jjj, :]-vrt_rhs[jjj, :])*scaleF, 'g-')
ax4.legend((p0[0], p1[0], p2[0]), \
(r"$\partial_t \overline{PV}^{z}$",\
r"$\sum RHS$", \
str(r"Residual ($\times$ %01.00e)" % (1/scaleF) )))
ax4.set_title(r"Residual -- X section")
plt.grid()
#- beta-plane -
llev1 = np.arange(-1.0, 1.1, 0.1)*5e-10
#
fig1 = plt.figure(figsize=(15, 10))
fig1.clf()
fig1.suptitle('beta-plane', fontsize=30)
#- rate -
ax1 = fig1.add_subplot(2, 3, 1)
cs1 = ax1.contourf(daVrt.vrt_rate[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax1.set_title(r'RATE')
#- wind stress -
ax2 = fig1.add_subplot(2, 3, 2)
cs2 = ax2.contourf(daVrt.vrt_wind[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax2.set_title(r'WIND STRESS')
#- planetary vorticity advection -
ax3 = fig1.add_subplot(2, 3, 3)
cs3 = ax3.contourf(daVrt.vrt_cor[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax3.set_title(r'Planetary vorticity advection')
#- NL advection -
ax4 = fig1.add_subplot(2, 3, 4)
cs4 = ax4.contourf(daVrt.vrt_xadv[ttt, :, :]+daVrt.vrt_yadv[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax4.set_title(r'NL advection')
#- horizontal mixing -
ax5 = fig1.add_subplot(2, 3, 5)
cs5 = ax5.contourf(daVrt.vrt_hmix[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax5.set_title(r'Hz dissipation')
#- bottom drag -
ax6 = fig1.add_subplot(2, 3, 6)
cs6 = ax6.contourf(daVrt.vrt_drag[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax6.set_title(r'BOTTOM DRAG')
#
cbax1 = fig1.add_axes([0.93, 0.25, 0.01, 0.5])
cb1 = fig1.colorbar(cs1, ax=ax1, orientation='vertical', cax=cbax1)
#
#str1 = ['rate', 'xadv', 'yadv', 'prsgrd', 'cor', 'vmix', 'hmix', 'wind', 'drag']
#for ip in range(9):
# exec( str("tmp=daVrt.vrt_%s[ttt, :, :]" % (str1[ip])) )
# ax = fig1.add_subplot(3, 3, ip+1)
# cs = ax.contourf(tmp, cmap='RdBu_r', levels=llev1, extend='both')
# plt.colorbar(cs)
# ax.set_title(str('vrt_%s' % str1[ip]))
# f-plane
llev1 = np.arange(-1.0, 1.1, 0.1)*5e-10
#
fig1 = plt.figure(figsize=(15, 10))
fig1.clf()
fig1.suptitle('f-plane', fontsize=30)
#- rate -
ax1 = fig1.add_subplot(2, 3, 1)
cs1 = ax1.contourf(daVrt2.vrt_rate[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax1.set_title(r'RATE')
#- wind stress -
ax2 = fig1.add_subplot(2, 3, 2)
cs2 = ax2.contourf(daVrt2.vrt_wind[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax2.set_title(r'WIND STRESS')
#- planetary vorticity advection -
ax3 = fig1.add_subplot(2, 3, 3)
cs3 = ax3.contourf(daVrt2.vrt_cor[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax3.set_title(r'Planetary vorticity advection')
#- NL advection -
ax4 = fig1.add_subplot(2, 3, 4)
cs4 = ax4.contourf(daVrt2.vrt_xadv[ttt, :, :]+daVrt2.vrt_yadv[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax4.set_title(r'NL advection')
#- horizontal mixing -
ax5 = fig1.add_subplot(2, 3, 5)
cs5 = ax5.contourf(daVrt2.vrt_hmix[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax5.set_title(r'Hz dissipation')
#- bottom drag -
ax6 = fig1.add_subplot(2, 3, 6)
cs6 = ax6.contourf(daVrt2.vrt_drag[ttt, :, :], cmap='RdBu_r', levels=llev1, extend='both')
ax6.set_title(r'BOTTOM DRAG')
#
cbax1 = fig1.add_axes([0.93, 0.25, 0.01, 0.5])
cb1 = fig1.colorbar(cs1, ax=ax1, orientation='vertical', cax=cbax1)