Post-processing-sensitivity/postpro_seas_icon_diff_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
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

### Saves the mean over a season from diff IRR-CTRL###

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

# It opens files, grid and extpar later
varname = "W_SO" # RELHUM_2M  'ishf' lhtfl t2m 'W_SO' # Choose the variable to work with

exp="24h_02" # 24h_01 "24h_02" #'24h_05'"sat03" #'fc06' # Choose exp

dr_data = "/hpc/uwork/extjroqu/out_bacy/BACYTEST_D03_ICONBC_"+str(exp)+"_nwp0/diff/"

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':
    dsets = xr.open_mfdataset(dr_data+"diff_*_"+str(varname)+".grb", combine='nested', concat_dim='time')
elif varname == "W_SO_4th":
    dsets = xr.open_mfdataset(dr_data+"diff_*_"+str(varname)+".grb", combine='nested', concat_dim='time')
elif varname == 'RELHUM_2M':
    dsets = xr.open_mfdataset(dr_data+"diff_*_RELHUM_2M.grb", combine='nested', concat_dim='time')
else:
    dsets = xr.open_mfdataset(dr_data+"diff_*_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 == "W_SO_4th":
    var = dsets["W_SO"]
elif varname == 'RELHUM_2M':
    var = dsets['r2']
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,"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,"Mean_"+str(exp)+"_"+str(varname)+"_"+str(aux_var.season.values)+"_seas.nc"))

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