plot_per_diff_seas_zoom_var.py

import numpy as np
from netCDF4 import Dataset
import xarray as xr
import pandas as pd
from glob import glob
import os
import os.path

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.tri as tri

import cartopy.crs as ccrs
import cartopy.feature as cfeature

from matplotlib import colors

# Plot seasonal mean for DIFF datasets, all variables from each experiment (fractional differences) for W_SO

# Output directory
save_to_path =os.path.abspath('/hpc/uwork/extjroqu/bacy_plots/p01_postprocessing/plots')

# It opens files, grid and extpar
exp="fc06" # "24h_02" #'24h_05' "24h_11" "sat03" #'fc06' # Choose exp
d_name = "Mean" #
labsize = 16

dr_data = "/hpc/uwork/extjroqu/bacy_plots/icon-nwp0_results/mean_data/"

dr_extra = "/hpc/uhome/extjroqu/bacy_data/"
gridfile = Dataset(dr_extra+"griddir/icon_grid_9999_R13B07_L.nc") #3km eu 
clon, clat  = np.rad2deg( gridfile.variables["clon"]) , np.rad2deg( gridfile.variables["clat"]) #[::-1] )

# Opening extpar dataset
extp = 'extpardir/icon_extpar_9999_R13B07_L_20200527_tiles.nc'
dext = xr.open_dataset(dr_extra+extp)
irri_crop = dext['LU_CLASS_FRACTION'][0,:]

indices_above_zero = np.where(irri_crop > 0)

o_land = dext["FR_LAND"]
o_land_mask = o_land > 0.5
o_land_mask = o_land_mask.rename({"cell": "values"})

###################### def plot  #######################
# Plot for different variables
def plot_icon (clon,clat,var):
    if varname == "W_SO":
        vnm =-100 
        vmx =95 
        levs =np.arange(vnm,vmx,5)
        cmap='RdBu'
    elif varname == "t2m":
        u = var.attrs['units']
        vnm = -2.5
        vmx = 2.75
        levs=np.arange(vnm,vmx,0.25)
        cmap='RdBu_r'
    elif varname == "lhtfl":
        u = var.attrs['units']
        vnm = -100 # -150
        vmx = 110 #165
        levs=np.arange(vnm,vmx,10)
        cmap='RdBu'
    elif varname == "RELHUM_2M":
        u = var.attrs['units']
        vnm = -35
        vmx = 38.5
        levs=np.arange(vnm,vmx,3.5)
        cmap='RdBu_r'
    else:
        u = var.attrs['units']
        vnm = -100
        vmx = 110
        levs=np.arange(vnm,vmx,10)
        cmap='RdBu'
    cmap_c = cmap
    levels = levs
    nlevs  = levels.size
    #norm = colors.TwoSlopeNorm(vmin=vnm, vcenter=0, vmax=vmx)
    projection = ccrs.PlateCarree()
# -- set colormap for nlevs
    cmap = plt.get_cmap(cmap_c, nlevs)
# -- create figure and axes instances; we need subplots for plot and colorbar
    fig, ax = plt.subplots(figsize=(8.7,7.5 ), # Width, height in inches
                       subplot_kw=dict(projection=projection))
# -- plot land areas at last to get rid of the contour lines at land
    gl = ax.gridlines(draw_labels=True,
                 linewidth=0.5,
                 color='dimgray',
                 alpha=0.4,
                 zorder=2)
    gl.xlabel_style = {'size': labsize, 'color': 'black'}
    gl.ylabel_style = {'size': labsize, 'color': 'black'}
    gl.top_labels = False
    gl.right_labels = False
    ax.coastlines(linewidth=0.5, zorder=2)
# -- contour plot
    cnf = ax.tricontourf(clon, clat, var,
                         levels=levels,
                         cmap=cmap,
                         extend='both',
                         norm=colors.CenteredNorm(),
                         zorder=0)
# -- If I want a zoom! ex. eu OR over Spain
    extent=[-10.9, 34.5, 29, 67.9] #lon, lat
    ax.set_extent(extent)         
# -- add a color bar
    cbar_ax = fig.add_axes([0.9, 0.25, 0.015, 0.5], autoscalex_on=True)
    cbar = fig.colorbar(cnf, cax=cbar_ax, orientation='vertical')
    cbar.ax.tick_params(axis='y', labelsize=labsize) 
    if varname == "t2m":
        cbar.ax.yaxis.set_major_formatter(plt.FormatStrFormatter('%.1f'))
    plt.setp(cbar.ax.get_xticklabels()[::2], visible=False)
    cbar.set_label("%",size=labsize)
    plt.savefig(os.path.join(save_to_path,"P100_DIFF_"+str(exp)+"_"+str(varname)+"_"+str(var.season.values)[0:10]+"_zoom.png"),dpi=300, bbox_inches='tight')
    plt.close()

####################### Mean for the whole area #######################
varname = "W_SO"

seas_list = ["MAM","JJA"]
for seasn in seas_list:
    dsets = xr.open_dataset(dr_data+str(d_name)+"_"+str(exp)+"_"+str(varname)+"_"+str(seasn)+"_per_seas.nc")
    print("Variable name in dsets ",dsets.data_vars)
    if varname == 'W_SO':
        var = dsets[varname] # For W_SO, time, levels, location
        var.coords #[I can use valid_time]
        var = var * 100 #Scale to %
    elif varname =='W_SO_4th':
        var = dsets['W_SO']
    else:
        var = dsets[varname]
    dsets.close()

    aux_var = var.where(o_land_mask)
    var_clean = aux_var.fillna(0)
    print("Shape aux_var ",var_clean)
    plot_icon(clon,clat,var_clean)

    print("Plot for DIFF mean for "+str(exp)+"_"+str(varname)+"_"+str(var.season.values))