plot_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

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

# It opens files, grid and extpar
exp="24h_02" # "24h_02" #'24h_05'"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 domain
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):
    labsize = 11

    if varname == "W_SO":
        u = var.units
        vnm = -2
        vmx = 2.25
        levs = np.arange(vnm,vmx,0.25)
        cmap='RdBu'
    elif varname == "W_SO_4th":
        u = var.units
        vnm = -6
        vmx = 6.75
        levs = np.arange(vnm,vmx,0.75)
        cmap='RdBu'
    elif varname == "t2m":
        u = var.attrs['units']
        vnm = -2.5
        vmx = 2.25
        levs=np.arange(vnm,vmx,0.25)
        cmap='RdBu_r'
    elif varname == "lhtfl":
        u = var.attrs['units']
        vnm = -100 # -150
        vmx = 90 #165
        levs=np.arange(vnm,vmx,10)
        cmap='RdBu'
    elif varname == "RELHUM_2M":
        u = var.attrs['units']
        vnm = -35
        vmx = 31.5
        levs=np.arange(vnm,vmx,3.5)
        cmap='RdBu_r'
    elif varname == 'Bowen_ratio':
        u = var.attrs['units']
        vnm = -5
        vmx = 5.8
        levs=np.arange(vnm,vmx,0.8)
        cmap='RdBu'
    else:
        u = var.attrs['units']
        vnm = -100
        vmx = 90
        levs=np.arange(vnm,vmx,10)
        cmap='RdBu'

    cmap_c = cmap
    levels = levs
    nlevs  = levels.size
    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(u,size=labsize)

    plt.savefig(os.path.join(save_to_path,"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 #######################
var_list = ["W_SO","t2m","lhtfl","ishf","RELHUM_2M"]

for v in var_list:
    varname = v

    seas_list = ["MAM","JJA"]
    for seasn in seas_list:
        dsets = xr.open_dataset(dr_data+str(d_name)+"_"+str(exp)+"_"+str(varname)+"_"+str(seasn)+"_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]
        elif varname =='W_SO_4th':
            var = dsets['W_SO']
        elif varname == 'RELHUM_2M':
            var = dsets['r2']
        elif varname == 'Bowen_ratio':
            var = dsets['ishf']
        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))