barakuda/python/exec/wind.py

#!/usr/bin/env python

#       B a r a K u d a
#
#     Generate global plots related to wind
#     wind stress, its curl, wind speed, etc
#     Use climatology fields built with 'build_clim.sh'
#
#       L. Brodeau, 2017

import sys
from os import path
import numpy as nmp
from netCDF4 import Dataset

import barakuda_tool as bt
import barakuda_plot as bp

cv_curl = 'socurl' ; # as generated by CDFTOOLS...
#CPAL_CURL='BrBG_r'
CPAL_CURL='ncview_jaisnb'

#CPAL_TAUM='ncview_nrl'
CPAL_TAUM='cubehelix_r'

CPAL_WNDM='CMRmap_r'

# Max values and increment for figures:
Rmax = 0.3 ; dR = 0.05  ; # Wind stress curl (10^-6 s)
Tmax = 0.3 ; dT = 0.025 ; # Wind stress module (N/m^2)
Wmax = 15. ; Wmin = 2. ; dW = 1.    ; # Surface wind speed module (m/s)
  
venv_needed = {'ORCA','EXP','DIAG_D','MM_FILE','FIG_FORM','FILE_FLX_SUFFIX','NN_TAUM','NN_WNDM'}

vdic = bt.check_env_var(sys.argv[0], venv_needed)

CONFEXP = vdic['ORCA']+'-'+vdic['EXP']

cd_clim = vdic['DIAG_D']+'/clim'

path_fig='./'
fig_type=vdic['FIG_FORM']


narg = len(sys.argv)
if narg < 3: print 'Usage: '+sys.argv[0]+' <year1> <year2>'; sys.exit(0)
cy1 = sys.argv[1] ; cy2=sys.argv[2]; jy1=int(cy1); jy2=int(cy2)
jy1_clim = jy1 ; jy2_clim = jy2
print ' => mean on the clim : ', jy1_clim, jy2_clim, '\n'

ctag = CONFEXP+'_'+cy1+'-'+cy2

# Getting coordinates:
bt.chck4f(vdic['MM_FILE'])
id_mm = Dataset(vdic['MM_FILE'])
xlon   = id_mm.variables['glamt'][0,:,:] ; xlat = id_mm.variables['gphit'][0,:,:]
Xmask = id_mm.variables['tmask'][0,0,:,:]
id_mm.close()

l_tau_is_annual = False
cf_nemo_curl = cd_clim+'/mclim_'+ctag+'_TCURL.nc4'

if not path.exists(cf_nemo_curl):
    cf_nemo_curl = cd_clim+'/aclim_'+ctag+'_TCURL.nc4'
    if path.exists(cf_nemo_curl):
        l_tau_is_annual = True
        print '\n *** wind.py : wind stress is annual...'
    else:
        print '\n *** WARNING: wind.py : giving up neither annual nor monthly curl clim found!'
        sys.exit(0)
    






#  Getting NEMO mean monthly climatology of wind stress module:
cextra_tau = ''
cv_taum = vdic['NN_TAUM']
if cv_taum != 'X':
    cf_nemo_taum = cd_clim+'/mclim_'+ctag+'_'+vdic['FILE_FLX_SUFFIX']+'.nc4'
    if l_tau_is_annual: cf_nemo_taum = cd_clim+'/aclim_'+ctag+'_'+vdic['FILE_FLX_SUFFIX']+'.nc4'
else:
    #Falling back on what cdfcurl.x has computed from monthly averaged Taux and Tauy:
    cf_nemo_taum = cf_nemo_curl
    cv_taum = 'sotaum'
    cextra_tau = ' (from monthly-averaged Tau_x & Tau_y !)'
    if l_tau_is_annual: cextra_tau = ' (from annually-averaged Tau_x & Tau_y !)'
    
bt.chck4f(cf_nemo_taum)
id_nemo = Dataset(cf_nemo_taum)
Xtaum   = id_nemo.variables[cv_taum][:,:,:]
id_nemo.close()
[ Nt, nj, ni ] = Xtaum.shape ; print ' Shape of TAUM :', Nt, nj, ni, '\n'

if not Nt in [1,12]:
    print '\n *** ERROR: wind.py : only accepting monthly or annual climatologies!', Nt
    sys.exit(0)



#  Getting NEMO mean monthly climatology of CURL:
bt.chck4f(cf_nemo_curl)
id_nemo = Dataset(cf_nemo_curl)
Xcurl   = id_nemo.variables[cv_curl][:,:,:]
id_nemo.close()

cextra_crl = ' (from monthly-averaged Tau_x & Tau_y !)'
if l_tau_is_annual: cextra_crl = ' (from annually-averaged Tau_x & Tau_y !)'


#  Getting surface wind speed as seen in NEMO if we find it!!!
l_do_wndm = False
cv_wndm = vdic['NN_WNDM']
if cv_wndm != 'X':
    cf_nemo_wndm = cd_clim+'/mclim_'+ctag+'_'+vdic['FILE_FLX_SUFFIX']+'.nc4'
    if path.exists(cf_nemo_wndm):        
        id_nemo = Dataset(cf_nemo_wndm)
        list_var = id_nemo.variables.keys()
        if cv_wndm in list_var:
            Xwndm   = id_nemo.variables[cv_wndm][:,:,:]
            l_do_wndm = True
            print '\n *** wind.py : we found wind speed in '+cf_nemo_wndm ; #lolo
        id_nemo.close()



nper = 3
if l_tau_is_annual: nper = 1

Xtaum_plot = nmp.zeros((nper,nj,ni))
Xcurl_plot = nmp.zeros((3,nj,ni))

Xtaum_plot[0,:,:] = nmp.mean(Xtaum[:,:,:]  ,axis=0) ;  # Annual
Xcurl_plot[0,:,:] = nmp.mean(Xcurl[:,:,:]  ,axis=0) ;  # Annual

if not l_tau_is_annual:
    Xtaum_plot[1,:,:] = nmp.mean(Xtaum[:3,:,:] ,axis=0) ;  # Winter
    Xtaum_plot[2,:,:] = nmp.mean(Xtaum[6:9,:,:],axis=0) ;  # Summer
    Xcurl_plot[1,:,:] = nmp.mean(Xcurl[:3,:,:] ,axis=0) ;  # Winter
    Xcurl_plot[2,:,:] = nmp.mean(Xcurl[6:9,:,:],axis=0) ;  # Summer


if l_do_wndm:
    Xwndm_plot = nmp.zeros((nper,nj,ni))
    Xwndm_plot[0,:,:] = nmp.mean(Xwndm[:,:,:]  ,axis=0) ;  # Annual
    Xwndm_plot[1,:,:] = nmp.mean(Xwndm[:3,:,:] ,axis=0) ;  # Winter
    Xwndm_plot[2,:,:] = nmp.mean(Xwndm[6:9,:,:],axis=0) ;  # Summer



# the Jean-Marc Molines method:
ji_lat0 = nmp.argmax(xlat[nj-1,:])

# Annual Curl:
bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xcurl_plot[0,:,:], Xmask, -Rmax, Rmax, dR,
              corca=vdic['ORCA'], lkcont=False, cpal=CPAL_CURL,
              cfignm=path_fig+'tau_curl_annual_'+CONFEXP, cbunit=r'$(10^{-6}s^{-1})$',
              ctitle='Wind stress curl, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_crl,
              lforce_lim=False, i_cb_subsamp=1,
              cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)

if not l_tau_is_annual:
    # JFM Curl:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xcurl_plot[1,:,:], Xmask, -Rmax, Rmax, dR,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_CURL,
                  cfignm=path_fig+'tau_curl_JFM_'+CONFEXP, cbunit=r'$(10^{-6}s^{-1})$',
                  ctitle='Wind stress curl, JFM, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_crl,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)

    # JAS Curl:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xcurl_plot[2,:,:], Xmask, -Rmax, Rmax, dR,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_CURL,
                  cfignm=path_fig+'tau_curl_JAS_'+CONFEXP, cbunit=r'$(10^{-6}s^{-1})$',
                  ctitle='Wind stress curl, JAS, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_crl,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)



# Annual Taum:
bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xtaum_plot[0,:,:], Xmask, 0., Tmax, dT,
              corca=vdic['ORCA'], lkcont=False, cpal=CPAL_TAUM,
              cfignm=path_fig+'taum_annual_'+CONFEXP, cbunit=r'$(N/m^2)$',
              ctitle='Wind stress module, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_tau,
              lforce_lim=False, i_cb_subsamp=1,
              cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)

if not l_tau_is_annual:
    # JFM Taum:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xtaum_plot[1,:,:], Xmask, 0., Tmax, dT,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_TAUM,
                  cfignm=path_fig+'taum_JFM_'+CONFEXP, cbunit=r'$(N/m^2)$',
                  ctitle='Wind stress module, JFM, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_tau,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)

    # JAS Taum:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xtaum_plot[2,:,:], Xmask, 0., Tmax, dT,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_TAUM,
                  cfignm=path_fig+'taum_JAS_'+CONFEXP, cbunit=r'$(N/m^2)$',
                  ctitle='Wind stress module, JAS, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_tau,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)



if l_do_wndm:
    # Annual Wndm:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xwndm_plot[0,:,:], Xmask, Wmin, Wmax, dW,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_WNDM,
                  cfignm=path_fig+'wndm_annual_'+CONFEXP, cbunit=r'$(m/s)$',
                  ctitle='Surface wind speed module, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_tau,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)
    
    # JFM Wndm:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xwndm_plot[1,:,:], Xmask, Wmin, Wmax, dW,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_WNDM,
                  cfignm=path_fig+'wndm_JFM_'+CONFEXP, cbunit=r'$(m/s)$',
                  ctitle='Surface wind speed module, JFM, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_tau,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)

    # JAS Wndm:
    bp.plot("2d")(xlon[0,:], xlat[:,ji_lat0], Xwndm_plot[2,:,:], Xmask, Wmin, Wmax, dW,
                  corca=vdic['ORCA'], lkcont=False, cpal=CPAL_WNDM,
                  cfignm=path_fig+'wndm_JAS_'+CONFEXP, cbunit=r'$(m/s)$',
                  ctitle='Surface wind speed module, JAS, '+CONFEXP+' ('+cy1+'-'+cy2+')'+cextra_tau,
                  lforce_lim=False, i_cb_subsamp=1,
                  cfig_type=fig_type, lat_min=-77., lat_max=75., lpix=True)