postpro_icon_plot_irr_area_SIM_std.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

# Plots all variables only for irrigated areas (SIMULATIONS)
# Output directory
save_to_path =os.path.abspath('/hpc/uwork/extjroqu/bacy_plots/p01_postprocessing/plots/daily_weight')

vn = "day"
timec = "daily"
country = "ALL" # ALL

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

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] )

####################### i_var has lon lat with irrigation grids #######################
var_list = ["W_SO","T_2M","LHFL_S","SHFL_S","RELHUM_2M"]
for v in var_list:
    # It opens files, grid and extpar
    varname = v

    dset2 = xr.open_dataset(dr_data+"SIM_mean_fc06_"+varname+"_day_irri.nc")
    dset1 = xr.open_dataset(dr_data+"SIM_mean_sat03_"+varname+"_day_irri.nc")
    dset3 = xr.open_dataset(dr_data+"SIM_mean_24h_11_"+varname+"_day_irri.nc")
    dset4 = xr.open_dataset(dr_data+"SIM_mean_24h_05_"+varname+"_day_irri.nc")
    dset5 = xr.open_dataset(dr_data+"SIM_mean_24h_02_"+varname+"_day_irri.nc")
    dset6 = xr.open_dataset(dr_data+"SIM_mean_ctrl_"+varname+"_day_irri.nc")

    ds_std2 = xr.open_dataset(dr_data+"SIM_std_fc06_"+varname+"_day_irri.nc")
    ds_std1 = xr.open_dataset(dr_data+"SIM_std_sat03_"+varname+"_day_irri.nc")
    ds_std3 = xr.open_dataset(dr_data+"SIM_std_24h_11_"+varname+"_day_irri.nc")
    ds_std4 = xr.open_dataset(dr_data+"SIM_std_24h_05_"+varname+"_day_irri.nc")
    ds_std5 = xr.open_dataset(dr_data+"SIM_std_24h_02_"+varname+"_day_irri.nc")
    ds_std6 = xr.open_dataset(dr_data+"SIM_std_ctrl_"+varname+"_day_irri.nc")

    exp_size = 6
    exp_coords = np.linspace(1, exp_size, exp_size)
    time_x = dset1.valid_time.size #if time is 14 weeks
    time_coords = dset1.valid_time.values  #if time is 14 weeks

    concat_mean = xr.zeros_like(xr.DataArray(np.empty([exp_size,time_x]), coords=([('exp', exp_coords),('time',time_coords)])))
    concat_mean['time'] = dset1.coords['valid_time'].values

    ds_list = [dset1, dset2, dset3, dset4, dset5, dset6]

    for d in range(exp_size):
        if varname == 'W_SO':
            concat_mean[d,:] = ds_list[d]["unknown"]
            concat_mean.attrs['units'] = dset1["unknown"].units
        elif varname == 'T_2M':
            concat_mean[d,:] = ds_list[d]["t2m"]
            concat_mean.attrs['units'] = dset1["t2m"].units
        elif varname == 'RELHUM_2M':
            concat_mean[d,:] = ds_list[d]["r2"]
            concat_mean.attrs['units'] = dset1["r2"].units
        elif varname == 'LHFL_S':
            concat_mean[d,:] = ds_list[d]["mslhf"]
            concat_mean.attrs['units'] = dset1["mslhf"].units
        elif varname == 'SHFL_S':
            concat_mean[d,:] = ds_list[d]["msshf"]
            concat_mean.attrs['units'] = dset1["msshf"].units
        else:
            print("No varname")

    df_w = pd.DataFrame(concat_mean, index=['SAT','FC','MIT','MSP','MFR','CTRL'], columns=concat_mean.time.values)
    
    concat_std = xr.zeros_like(xr.DataArray(np.empty([exp_size,time_x]), coords=([('exp', exp_coords),('time',time_coords)])))
    concat_std['time'] = dset1.coords['valid_time'].values

    ds_list = [ds_std1, ds_std2, ds_std3, ds_std4, ds_std5, ds_std6]

    for d in range(exp_size):
        if varname == 'W_SO':
            concat_std[d,:] = ds_list[d]["unknown"]
            concat_std.attrs['units'] = dset1["unknown"].units
        elif varname == 'T_2M':
            concat_std[d,:] = ds_list[d]["t2m"]
            concat_std.attrs['units'] = dset1["t2m"].units
        elif varname == 'RELHUM_2M':
            concat_std[d,:] = ds_list[d]["r2"]
            concat_std.attrs['units'] = dset1["r2"].units
        elif varname == 'LHFL_S':
            concat_std[d,:] = ds_list[d]["mslhf"]
            concat_std.attrs['units'] = dset1["mslhf"].units
        elif varname == 'SHFL_S':
            concat_std[d,:] = ds_list[d]["msshf"]
            concat_std.attrs['units'] = dset1["msshf"].units
        else:
            print("No varname")

    df_std = pd.DataFrame(concat_std, index=['SAT','FC','MIT','MSP','MFR','CTRL'], columns=concat_std.time.values)


    # Plot the data for other variables
    df_w_transposed = df_w.T
    df_w_std_transposed = df_std.T
    # Specify custom colors, second option colorblind friendly
    #custom_colors = ['red', 'blue','purple', 'green', 'orange']
    custom_colors = ["#D55E00", "#0072B2","#984ea3" ,"#009E73", "#E69F00"]

    fig, ax = plt.subplots(figsize=(10, 4))
    for column, color in zip(df_w_transposed.columns, custom_colors):
        plt.plot(df_w_transposed.index, df_w_transposed[column], label=column, color=color)

    for column, color in zip(df_w_std_transposed.columns, custom_colors):
        ax.fill_between(df_w_std_transposed.index,
                    df_w_transposed[column] - df_w_std_transposed[column],
                    df_w_transposed[column] + df_w_std_transposed[column],
                    color=color, alpha=0.2)
    if varname == "W_SO" or varname == "RELHUM_2M":
        plt.legend(loc='lower left',bbox_to_anchor =(0, 0))
        plt.grid()
        plt.xticks(rotation=45, ha='right', fontsize=12)
        plt.ylabel("\u0394 "+str(varname)+" ("+str(concat_mean.units)+")", fontsize=12)
    else:
        fontlab = 16
        plt.legend(loc='lower left',bbox_to_anchor =(0, 0), fontsize=12)
        plt.grid()
        plt.xticks(rotation=45, ha='right', fontsize=fontlab)
        plt.ylabel("\u0394 "+str(varname)+" ("+str(concat_mean.units)+")", fontsize=fontlab)
    plt.tight_layout() #Remove excess of white space

    plt.savefig(os.path.join(save_to_path,"SIM_merged_"+str(varname)+"_"+str(timec)+"_irri_std.png"),dpi=300)

    print("DONE plot irr mean for all exp "+str(timec)+" "+str(varname))

    print("Variable MAX  for "+str(varname)+": ",concat_mean.max(dim='time').values)
    print("Variable MIN  for "+str(varname)+": ",concat_mean.min(dim='time').values)
    print("Variable MEAN  for "+str(varname)+": ",concat_mean.mean(dim='time').values)

print("DONE with all variables!")