-
Christoph.Fischer authoredChristoph.Fischer authored
climatology.py 3.86 KiB
import numpy as np
import xarray as xr
import metpy.calc as mpcalc
def smthClmDay(clmDay, nHarm): # from athul
# 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)
# also weird dims. le_ti_level_0
def compute_climatology(cfg):
strt = '1979-01-01'
last = '2019-12-31'
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),
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