Post-processing-sensitivity/postpro_seas_icon_simu_save.py at main · jlroquem/Post-processing-sensitivity · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
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

# Output directory
save_to_path =os.path.abspath('/hpc/uwork/extjroqu/bacy_plots/icon-nwp0_results/mean_data/')

# It opens files, grid and extpar
varname = "RELHUM_2M" #  W_SO.4th RELHUM_2M  'ishf' lhtfl t2m 'W_SO' # Choose the variable to work with
exp="24h_11" # ctrl 24h_11 "24h_02" #'24h_05'"sat03" #'fc06' # Choose exp
i_ver="nwp0"

if varname == "W_SO":
    dr_data = "/hpc/uwork/extjroqu/BACYTEST_D03_ICONBC_"+str(exp)+"_"+str(i_ver)+"/test_1/"
elif varname == 'RELHUM_2M':
    dr_data = "/hpc/uwork/extjroqu/BACYTEST_D03_ICONBC_"+str(exp)+"_"+str(i_ver)+"/test_4/"
elif varname == 'bowen':
    dr_data = "/hpc/uwork/extjroqu/BACYTEST_D03_ICONBC_"+str(exp)+"_"+str(i_ver)+"/test_6/"
elif varname == "W_SO.4th":
    dr_data = "/hpc/uwork/extjroqu/BACYTEST_D03_ICONBC_"+str(exp)+"_"+str(i_ver)+"/test_8/"
else:
    dr_data = "/hpc/uwork/extjroqu/BACYTEST_D03_ICONBC_"+str(exp)+"_"+str(i_ver)+"/test_2/"

print("The path from where data comes from is ",dr_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] )


if varname == 'W_SO': # exp new exp and fc the old one
    dsets = xr.open_mfdataset(dr_data+"fc_DOM01_*.grb."+str(varname)+".grb", combine='nested', concat_dim='time')
elif varname == "W_SO.4th":
    dsets = xr.open_mfdataset(dr_data+"fc_DOM01_*.grb."+str(varname)+".grb", combine='nested', concat_dim='time')
elif varname == 'RELHUM_2M':
    dsets = xr.open_mfdataset(dr_data+"fc_DOM01_*.grb.RELHUM_2M.grb", combine='nested', concat_dim='time')
elif varname == 'bowen':
    dsets = xr.open_mfdataset(dr_data+"fc_DOM01_*_"+str(varname)+".grb", combine='nested', concat_dim='time')
else:
    dsets = xr.open_mfdataset(dr_data+"fc_DOM01_*.grb.3var.grb", combine='nested', concat_dim='time')

print("Variable name in dsets ",dsets.data_vars)

# 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)
print("Indices of values above zero:", indices_above_zero)

if varname == 'W_SO':
    var = dsets[varname] # For W_SO, time, levels, location
    var.coords #[I can use valid_time]
elif varname == 'RELHUM_2M':
    var = dsets['r2']
elif varname == 'bowen':
    var = dsets['ishf']
elif varname == "W_SO.4th":
    var = dsets["W_SO"]
else:
    var = dsets[varname]
    print("Variable attributes ",var.attrs)

####################### Mean for the whole area #######################
# Release some memory that I do not need anymore
dsets.close()

x_all = var
x_all['time'] = x_all.coords['valid_time'].values

if varname == 'W_SO':
    # Mean over layers and then mean over a season
    jja_lev_mean_all = x_all.mean(dim='depthBelowLandLayer').groupby('time.season').mean()
    jja_lev_mean_all.attrs['units'] = x_all.units
    jja_lev_mean_all.attrs['GRIB_shortName'] = x_all.GRIB_shortName
    print("Number of iterations over time ",jja_lev_mean_all.shape[0])
    #Creating an aux var for each level
    for t in range(jja_lev_mean_all.shape[0]):
        aux_var = jja_lev_mean_all[t,:].compute()
        aux_var[np.isnan(aux_var)] = 0
        print("aux_var shape for each iteration ",aux_var.shape)
        aux_var.to_netcdf(os.path.join(save_to_path,"SIM_Mean_"+str(exp)+"_"+str(varname)+"_"+str(aux_var.season.values)+"_seas.nc"))
else:
    # mean over a season for the other variables
    sea_mean_all = x_all.groupby('time.season').mean()
    for t in range(sea_mean_all.shape[0]):
        aux_var = sea_mean_all[t,:].compute()
        aux_var[np.isnan(aux_var)] = 0
        aux_var.to_netcdf(os.path.join(save_to_path,"SIM_Mean_"+str(exp)+"_"+str(varname)+"_"+str(aux_var.season.values)+"_seas.nc"))

print("Save dataset over season for "+str(exp)+"_"+str(varname)+" DONE!")