source: NEMO/branches/2020/dev_r13648_ASINTER-04_laurent_bulk_ice/tests/STATION_ASF/EXPREF/plot_station_asf_OI.py @ 13690

#!/usr/bin/env python3
# -*- Mode: Python; coding: utf-8; indent-tabs-mode: nil; tab-width: 4 -*-
#
##########################################################################################
# Post-diagnostic of STATION_ASF with sea-ice support (ex: run with forcing "ERA5_arctic")
#
#  L. Brodeau, 2020
##########################################################################################

import sys
from os import path as path
import math
import numpy as nmp
from netCDF4 import Dataset,num2date
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

CONFIG='STATION_ASF'

cforcing = 'PAPA' ; # name of forcing ('PAPA', 'ERA5_arctic', etc...)

# (files are: output/3x3/'+CONFIG+'-'+algo+'_1h_'+year+'0101_'+year+'1231_gridT_'+cforcing+'.nc' )
cstation = 'ERA5 81N, 36.75E'

cy1     = '2018' ; # First year
cy2     = '2018' ; # Last year

jt0 = 0

dir_figs='.'
size_fig=(13,8)
size_fig0=(12,10)
fig_ext='png'

clr_red = '#AD0000'
clr_sat = '#ffed00'
clr_mod = '#008ab8'

rDPI=100.

L_ALGOS = [ 'ECMWF-LG15', 'ECMWF-LU12', 'ECMWF-CSTC' ]
l_color = [  '#ffed00'  , '#008ab8'   ,     '0.4'    ] ; # colors to differentiate algos on the plot
l_width = [     3       ,      2      ,       1      ] ; # line-width to differentiate algos on the plot
l_style = [    '-'      ,     '-'     ,      '--'    ] ; # line-style


# Variables to compare for A GIVEN algorithm
###############################################
#L_VNEM0 = [ 



# Variables to compare between algorithms
############################################
L_VNEM = [   'Cd_ice',   'Ce_ice',   'qla_ice' ,   'qsb_ice'   ,  'qt_ice'    ,  'qlw_ice'  ,  'qsr_ice'  , 'taum_ai'   ]
L_VARO = [     'Cd'  ,     'Ce'  ,   'Qlat'    ,    'Qsen'     ,   'Qnet'     ,   'Qlw'     ,    'Qsw'    ,  'Tau'      ]
L_VARL = [ r'$C_{D}$', r'$C_{E}$', r'$Q_{lat}$', r'$Q_{sens}$' , r'$Q_{net}$' , r'$Q_{lw}$' , r'$Q_{sw}$' , r'$|\tau|$' ]
L_VUNT = [     ''    ,     ''    , r'$W/m^2$'  , r'$W/m^2$'    , r'$W/m^2$'   , r'$W/m^2$'  , r'$W/m^2$'  , r'$N/m^2$'  ]
L_VMAX = [    0.003  ,    0.003  ,     75.     ,     75.       ,    800.      ,    200.     ,    200.     ,    1.2      ]
L_VMIN = [    0.     ,    0.     ,   -250.     ,   -125.       ,   -400.      ,   -200.     ,      0.     ,    0.       ]
L_ANOM = [   False   ,   False   ,   True      ,    True       ,    True      ,    True     ,    True     ,   True      ]


nb_algos = len(L_ALGOS)

# Getting arguments:
narg = len(sys.argv)
if narg != 2:
    print('Usage: '+sys.argv[0]+' <DIR_OUT_SASF>'); sys.exit(0)
cdir_data = sys.argv[1]



# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
# Populating and checking existence of files to be read
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
def chck4f(cf):
    cmesg = 'ERROR: File '+cf+' does not exist !!!'
    if not path.exists(cf): print(cmesg) ; sys.exit(0)

cf_in = []
for ja in range(nb_algos):
    cfi = cdir_data+'/output/3x3/'+CONFIG+'-'+L_ALGOS[ja]+'_1h_'+cy1+'0101_'+cy2+'1231_icemod_'+cforcing+'.nc'
    chck4f(cfi)
    cf_in.append(cfi)
print('Files we are goin to use:')
for ja in range(nb_algos): print(cf_in[ja])
#-----------------------------------------------------------------

# Getting time array from the first file:
id_in = Dataset(cf_in[0])
vt = id_in.variables['time_counter'][jt0:]
cunit_t = id_in.variables['time_counter'].units ; print(' "time_counter" is in "'+cunit_t+'"')
id_in.close()
Nt = len(vt)

vtime = num2date(vt, units=cunit_t) ; # something understandable!
vtime = vtime.astype(dtype='datetime64[D]')

ii=Nt/300
ib=max(ii-ii%10,1)
xticks_d=int(30*ib)

rat = 100./float(rDPI)
params = { 'font.family':'Open Sans',
           'font.size':       int(15.*rat),
           'legend.fontsize': int(15.*rat),
           'xtick.labelsize': int(15.*rat),
           'ytick.labelsize': int(15.*rat),
           'axes.labelsize':  int(16.*rat)
}
mpl.rcParams.update(params)
font_inf = { 'fontname':'Open Sans', 'fontweight':'normal', 'fontsize':18.*rat }
font_x   = { 'fontname':'Open Sans', 'fontweight':'normal', 'fontsize':15.*rat }


# Now we compare output variables from bulk algorithms between them:

nb_var = len(L_VNEM)

xF  = nmp.zeros((Nt,nb_algos))
xFa = nmp.zeros((Nt,nb_algos))


for jv in range(nb_var):
    print('\n *** Treating variable: '+L_VARO[jv]+' !')

    for ja in range(nb_algos):
        #
        id_in = Dataset(cf_in[ja])
        xF[:,ja] = id_in.variables[L_VNEM[jv]][jt0:,1,1] # only the center point of the 3x3 spatial domain!
        if ja == 0: cvar_lnm = id_in.variables[L_VNEM[jv]].long_name
        id_in.close()

    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    fig = plt.figure(num = jv, figsize=size_fig, facecolor='w', edgecolor='k')
    ax1 = plt.axes([0.08, 0.25, 0.9, 0.7])
    ax1.set_xticks(vtime[::xticks_d])
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M:%S'))
    plt.xticks(rotation='60', **font_x)

    for ja in range(nb_algos):
        plt.plot(vtime, xF[:,ja], '-', color=l_color[ja], linestyle=l_style[ja], linewidth=l_width[ja], label=L_ALGOS[ja], zorder=10+ja)

    ax1.set_ylim(L_VMIN[jv], L_VMAX[jv]) ; ax1.set_xlim(vtime[0],vtime[Nt-1])
    plt.ylabel(L_VARL[jv]+' ['+L_VUNT[jv]+']')

    ax1.grid(color='k', linestyle='-', linewidth=0.3)
    plt.legend(loc='best', ncol=1, shadow=True, fancybox=True)
    ax1.annotate(cvar_lnm+', station: '+cstation, xy=(0.3, 1.), xycoords='axes fraction',  bbox={'facecolor':'w', 'alpha':1., 'pad':10}, zorder=50, **font_inf)
    plt.savefig(L_VARO[jv]+'.'+fig_ext, dpi=int(rDPI), transparent=False)
    plt.close(jv)
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    def symetric_range( pmin, pmax ):
        # Returns a symetric f-range that makes sense for the anomaly of "f" we're looking at...
        from math import floor, copysign, log, ceil
        zmax = max( abs(pmax) , abs(pmin) )
        romagn = floor(log(zmax, 10)) ; # order of magnitude of the anomaly  we're dealing with
        rmlt = 10.**(int(romagn)) / 2.
        frng = copysign( ceil(abs(zmax)/rmlt)*rmlt , zmax)
        return frng

    

    if L_ANOM[jv]:

        for ja in range(nb_algos): xFa[:,ja] = xF[:,ja] - nmp.mean(xF,axis=1)

        if nmp.sum(nmp.abs(xFa[:,:])) == 0.0:
            print('     Well! Seems that for variable '+L_VARO[jv]+', choice of algo has no impact a all!')
            print('          ==> skipping anomaly plot...')

        else:

            yrng = symetric_range( nmp.min(xFa) , nmp.max(xFa) )

            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
            fig = plt.figure(num = 10+jv, figsize=size_fig, facecolor='w', edgecolor='k')
            ax1 = plt.axes([0.08, 0.25, 0.9, 0.7])
            ax1.set_xticks(vtime[::xticks_d])
            ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M:%S'))
            plt.xticks(rotation='60', **font_x)

            for ja in range(nb_algos):
                plt.plot(vtime, xFa[:,ja], '-', color=l_color[ja], linewidth=l_width[ja], label=L_ALGOS[ja], zorder=10+ja)

            ax1.set_ylim(-yrng,yrng) ; ax1.set_xlim(vtime[0],vtime[Nt-1])
            plt.ylabel(L_VARL[jv]+' ['+L_VUNT[jv]+']')
            ax1.grid(color='k', linestyle='-', linewidth=0.3)
            plt.legend(bbox_to_anchor=(0.45, 0.2), ncol=1, shadow=True, fancybox=True)
            ax1.annotate('Anomaly of '+cvar_lnm, xy=(0.3, 0.97), xycoords='axes fraction',  bbox={'facecolor':'w', 'alpha':1., 'pad':10}, zorder=50, **font_inf)
            plt.savefig(L_VARO[jv]+'_anomaly.'+fig_ext, dpi=int(rDPI), transparent=False)
            plt.close(10+jv)
            #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~