高质量代码计算降水季节指数Seasonality index

以CMIP6模式为例，计算结果可复现，检验正确

import xarray as xr
da = xr.open_dataset(r'/home/mw/input/cmip7374/pr_Amon_ACCESS-ESM1-5_historical_ensemble_05x05_190101-201412.nc')
pre = da.pr* 60 * 60 * 24
import xarray as xr
import numpy as np
lon = pre.lon[(pre.lon>=70)&(pre.lon<=140)]
lat = pre.lat[(pre.lat>=15)&(pre.lat<=60)]
print(lat)
pre = pre.sel(lat=lat,lon=lon)
k = pre.time.dt.year
kyear = len(list(k.groupby('time.year').groups.keys()))
r_avg_m = pre.groupby('time.month').sum()/kyear
R_avg = r_avg_m.sum('month')
R_avg.max()
p_avg_m = r_avg_m/R_avg 
D_avg = p_avg_m*np.log2(p_avg_m/(1./12.))
D = D_avg.sum('month')
SI = D*R_avg/R_avg.max()

m,p = xr.broadcast(pre.time.dt.month,pre)
mpre = m*pre
r_k_m = pre.groupby('time.year')
C_k = mpre.groupby('time.year').sum()/pre.groupby('time.year').sum()




mon_m = np.array(m.values).reshape(len(m.time)//12,12,m.shape[1],m.shape[2])
r_k_m = np.array(pre.values).reshape(len(pre.time)//12,12,pre.shape[1],pre.shape[2])
sumval = r_k_m*(mon_m-np.broadcast_to(C_k.expand_dims(dim={"mon": 1},axis=1),mon_m.shape))**2
R_k = pre.groupby('time.year').sum()
Z_k = np.sqrt(sumval.sum(axis=1)/R_k)

矢量化计算nc日值的数据里，求每年所有日值中最大的两个值之和，适用于cmip数据集计算

如果还想再快一点可以使用map_block或dask加速

import xarray as xr
import heapq

def vectorize(var):
    #use numpy original sound
    def find_max(value):
        # find max top2 values
        return np.nansum(heapq.nlargest(2,value))
    
    return xr.apply_ufunc(find_max,var,input_core_dims=[['time']],output_core_dims = [[]],vectorize=True)
    
    
def aggfunc(var):
    aggvar = vectorize(var)
    return aggvar

postdata = dasel.pr.groupby(dasel.time.dt.year).map(aggfunc)