Newer
Older
Christoph Fischer
committed
import numpy as np
import xarray as xr
import metpy.calc as mpcalc
Christoph Fischer
committed
# nharm is the number of harmonic modes that needs to be removed. So we retain that in smthClm to later
# subtract from raw.
print(clmDay)
cf = np.fft.rfft(clmDay, axis=0) # ([nday/2]+1)x[nlat]x[mlon]<==n:even
# ([nday+1/2])x[nlat]x[mlon]<==n:odd
# cf(:,0:nharm-1)
cf[nHarm, :, :] = 0.5 * cf[nHarm, :, :] # mini-taper
cf[nHarm + 1:, :, :] = 0.0 # set all higher coef. to zero
clmDaySmth = np.fft.irfft(cf, axis=0) # reverse FT to real reconstruct data
clmDaySmth = xr.DataArray(clmDaySmth, dims=('hourofyear', 'lat', 'lon'),
coords={'hourofyear': clmDay.hourofyear, 'lat': clmDay.lat, 'lon': clmDay.lon})
return (clmDaySmth)
# TODO TEST THIS, maybe needs fix, currently cached! (maybe works with multi year data)
# also weird dims. le_ti_level_0
Christoph Fischer
committed
def compute_climatology(cfg):
strt = '1979-01-01'
last = '2019-12-31'
Christoph Fischer
committed
cv_database_dir = cfg.cv_data_dir
f = xr.open_mfdataset(cv_database_dir + "*.nc").sel(
# plev=slice(plev, plev),
# lat=slice(latS,latN),
# lon=slice(lonW,lonE),
Christoph Fischer
committed
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
time=slice(strt, last))
# smoothing
cv = mpcalc.smooth_n_point(f.cv, n=9, passes=2).metpy.dequantify()
# generate hourofyear climatology, save it to file for next time
# Calculate a hour-of-year climatology to be subtracted later.
cv = cv.assign_coords(hourofyear=cv.time.dt.strftime("%m-%d %H"))
print('Computing climatology...')
cv_clim_da = cv.groupby('hourofyear').mean('time')
# print('Smoothing climatology...')
# clm = smthClmDay(clm, 3) # calculate climatology with just the 3 harmonic modes retained.
print('Climatology computed ')
cv_clim_da.attrs['units'] = 's**-1'
cv_clim_da.attrs['standard_name'] = 'Curvature Vorticity'
cv_clim_da.attrs['long_name'] = 'Curvature Vorticity climatology (hour of year)'
cv_clim = cv_clim_da.to_dataset(name="cv")
# compute anomalies
print("Compute anomalies of " + str(cfg.cv_percentile) + " percentile (total)...")
full_cv_anomalies = cv.groupby('hourofyear') - cv_clim.cv
# quantile applied over time,lat,lon -> so one scalar of cfg.cv_percentile% quantile per pressure level
quantiles_full = full_cv_anomalies.quantile(float(cfg.cv_percentile) / 100.0, dim=['lat', 'lon', 'time'])
cv_clim['cva_quantile_full'] = quantiles_full
print(quantiles_full.data)
print("Compute anomalies of " + str(cfg.cv_percentile) + " percentile (hourofyear)...")
# quantile for each hourofyear, bit more complicated. bascially for each Jan 1 00UTC the x% quantile etc.
full_cv_anomalies = full_cv_anomalies.assign_coords(hourofyear=full_cv_anomalies.time.dt.strftime("%m-%d %H"))
full_cv_anomalies = full_cv_anomalies.swap_dims({"time": "hourofyear"})
full_cv_anomalies_stacked = full_cv_anomalies.stack(le_ti=("plev", "hourofyear"))
# group by pressure level and hourofyear -> quantile over lat+lon+date(hoy)
quantiles_hoy_stacked = full_cv_anomalies_stacked.groupby('le_ti').quantile(q=float(cfg.cv_percentile) / 100.0,
dim=...)
quantiles_hoy = quantiles_hoy_stacked.unstack(dim='le_ti')
cv_clim['cva_quantile_hoy'] = quantiles_hoy
quantiles_hoy = cv_clim.cva_quantile_hoy
# somehow xarray cant reconstruct stacked groupby results, fix manually.
quantiles_hoy = quantiles_hoy.swap_dims({"le_ti_level_0": "plev"})
quantiles_hoy = quantiles_hoy.swap_dims({"le_ti_level_1": "hourofyear"})
quantiles_hoy = quantiles_hoy.rename({"le_ti_level_0": "plev"})
quantiles_hoy = quantiles_hoy.rename({"le_ti_level_1": "hourofyear"})
print(quantiles_hoy)
cv_clim['cva_quantile_hoy'] = quantiles_hoy
cv_clim = cv_clim.drop_dims(["le_ti_level_0", "le_ti_level_1"])
print("Done!")
print("Generated dataset:")
print(cv_clim)
return cv_clim