利用Dask提高气候变化下全球光计算的速度

Question

我正在实现一种估计海洋表面光的算法，其函数是风(波浪、表面粗糙度)、叶绿素和天顶角。我想要这样做，使用CMIP6的气候预测作为1950年至2100年期间每月的投入。我使用Python和木星笔记本来读取来自谷歌云提供的CMIP6 6气候模型的云、叶绿素和风的全局值。

完整代码是这里提供的木星笔记本电脑.

我使用Python库 pvlib计算海洋表面的直接和漫射光，这取决于来自CMIP6模型的时间、地理位置和云层。我用Seferian等人。2018年在同一时间和地点计算叶绿素和波浪估计光的反照率的方法。我的代码中的瓶颈似乎是估计def calculate_OSA函数中的波和叶绿素对光反照率的影响，该函数估计10 nm间隔的所有波长200-4000 nm的反射光谱。我使用numpy vectorized循环给定地理网格点的所有波长，并使用dask.delayed循环所有网格点。全球覆盖的Gridpoint为180x360。

def calculate_OSA(µ_deg, uv, chl, wavelengths, refractive_indexes, alpha_chl, alpha_w, beta_w, alpha_wc, solar_energy):
    if (µ_deg<0 or µ_deg>180):
        µ_deg=0

        µ = np.cos(np.radians(µ_deg))

        # Solar zenith angle
        # wind is wind at 10 m height (m/s)
        σ = np.sqrt(0.003+0.00512*uv)

        # Vectorize the functions
        vec_calculate_direct_reflection=np.vectorize(calculate_direct_reflection)
        vec_calculate_diffuse_reflection=np.vectorize(calculate_diffuse_reflection)
        vec_calculate_direct_reflection_from_chl=np.vectorize(calculate_direct_reflection_from_chl)
        vec_calculate_diffuse_reflection_from_chl=np.vectorize(calculate_diffuse_reflection_from_chl)

        # Direct reflection
        alpha_direct = vec_calculate_direct_reflection(refractive_indexes,µ,σ)

        # Diffuse reflection
        alpha_diffuse = vec_calculate_diffuse_reflection(refractive_indexes,σ)

        # Reflection from chlorophyll and biological pigments
        alpha_direct_chl = vec_calculate_direct_reflection_from_chl(wavelengths, chl, alpha_chl, alpha_w, beta_w, σ, µ, alpha_direct)

        # Diffuse reflection interior of water from chlorophyll
        alpha_diffuse_chl = vec_calculate_diffuse_reflection_from_chl(wavelengths, chl, alpha_chl, alpha_w, beta_w, σ, alpha_direct)

        # OSA
        return 

整个脚本被写成一个朱佩笔记本，尽管它使用一个子程序来读取CMIP6数据和一个反照率计算记事本。我知道脚本是长而复杂的，但我认为可以改进的主要功能是def calculate_OSA和主calculate_light函数。在calculate_light中，我相信我可以改进如何使用dask和分块，以及如何将calculate_light中的主循环矢量化可以加快速度。

目前，在带有16 of内存的mini上运行一个时间步骤需要2.27分钟。

%%time
def calculate_light(config_pices_obj):

    selected_time=0
    wavelengths, refractive_indexes, alpha_chl, alpha_w, beta_w, alpha_wc, solar_energy = albedo.setup_parameters()
    startdate=datetime.datetime.now()

    regional=True
    create_plots=True

    southern_limit_latitude=45
    for key in config_pices_obj.dset_dict.keys():

        var_name = key.split("_")[0]
        model_name = key.split("_")[3]

        if var_name=="uas":

            key_v="vas"+key[3:]
            key_chl="chl"+key[3:]
            key_clt="clt"+key[3:]
            key_sisnconc="sisnconc"+key[3:]
            key_sisnthick="sisnthick"+key[3:]
            key_siconc="siconc"+key[3:]
            key_sithick="sithick"+key[3:]

            var_name_v = key_v.split("_")[0]
            model_name_v = key_v.split("_")[3]

            print("=> model: {} variable name: {}".format(key, var_name))
            print("=> model: {} variable name: {}".format(key_v, var_name_v))

            if model_name_v==model_name:

                ds_uas=config_pices_obj.dset_dict[key].isel(time=selected_time) 
                ds_vas=config_pices_obj.dset_dict[key_v].isel(time=selected_time)
                ds_chl=config_pices_obj.dset_dict[key_chl].isel(time=selected_time)
                ds_clt=config_pices_obj.dset_dict[key_clt].isel(time=selected_time)
                ds_sisnconc=config_pices_obj.dset_dict[key_sisnconc].isel(time=selected_time)
                ds_sisnthick=config_pices_obj.dset_dict[key_sisnthick].isel(time=selected_time)
                ds_siconc=config_pices_obj.dset_dict[key_siconc].isel(time=selected_time)
                ds_sithick=config_pices_obj.dset_dict[key_sithick].isel(time=selected_time)

                if regional:
                    ds_uas=ds_uas.sel(y=slice(southern_limit_latitude,90))
                    ds_vas=ds_vas.sel(y=slice(southern_limit_latitude,90))
                    ds_chl=ds_chl.sel(y=slice(southern_limit_latitude,90))
                    ds_clt=ds_clt.sel(y=slice(southern_limit_latitude,90))
                    ds_sisnconc=ds_sisnconc.sel(y=slice(southern_limit_latitude,90))
                    ds_sisnthick=ds_sisnthick.sel(y=slice(southern_limit_latitude,90))
                    ds_siconc=ds_siconc.sel(y=slice(southern_limit_latitude,90))
                    ds_sithick=ds_sithick.sel(y=slice(southern_limit_latitude,90))

                # Regrid to cartesian grid:
                # For any Amon related variables (wind, clouds), the resolution from CMIP6 models is less than 
                # 1 degree longitude x latitude. To interpolate to a 1x1 degree grid we therefore first interpolate to a 
                # 2x2 degrees grid and then subsequently to a 1x1 degree grid.

                ds_out_amon = xe.util.grid_2d(-180,180,2,southern_limit_latitude,90,2) 
                ds_out = xe.util.grid_2d(-180,180,1,southern_limit_latitude,90,1) #grid_global(1, 1)

                dr_out_uas_amon=regrid_variable("uas",ds_uas,ds_out_amon,transpose=True).to_dataset()
                dr_out_uas=regrid_variable("uas",dr_out_uas_amon,ds_out,transpose=False)

                dr_out_vas_amon=regrid_variable("vas",ds_vas,ds_out_amon,transpose=True).to_dataset()
                dr_out_vas=regrid_variable("vas",dr_out_vas_amon,ds_out,transpose=False)

                dr_out_clt_amon=regrid_variable("clt",ds_clt,ds_out_amon,transpose=True).to_dataset()
                dr_out_clt=regrid_variable("clt",dr_out_clt_amon,ds_out,transpose=False)
                dr_out_chl=regrid_variable("chl",ds_chl,ds_out,transpose=False)

                dr_out_sisnconc=regrid_variable("sisnconc",ds_sisnconc,ds_out,transpose=False)
                dr_out_sisnthick=regrid_variable("sisnthick",ds_sisnthick,ds_out,transpose=False)
                dr_out_siconc=regrid_variable("siconc",ds_siconc,ds_out,transpose=False)
                dr_out_sithick=regrid_variable("sithick",ds_sithick,ds_out,transpose=False)

                # Calculate scalar wind and organize the data arrays to be used for  given timestep (month-year)
                wind=np.sqrt(dr_out_uas**2+dr_out_vas**2).values

                lat=dr_out_uas.lat.values
                lon=dr_out_uas.lon.values

                clt=dr_out_clt.values
                chl=dr_out_chl.values
                sisnconc=dr_out_sisnconc.values
                sisnthick=dr_out_sisnthick.values
                siconc=dr_out_siconc.values
                sithick=dr_out_sithick.values

                m=len(wind[:,0])
                n=len(wind[0,:])
                month=6

                all_direct=[]
                all_OSA=[]
                for hour_of_day in range(12,13,1):
                    print("Running for hour {}".format(hour_of_day))

                    calc_radiation = [dask.delayed(radiation)(clt[j,:],lat[j,0],month,hour_of_day) for j in range(m)]

                    # https://github.com/dask/dask/issues/5464   
                    rad = dask.compute(calc_radiation, scheduler='processes')
                    rads=np.asarray(rad).reshape((m, n, 3))

                    zr = [dask.delayed(calculate_OSA)(rads[i,j,2], wind[i,j], chl[i,j], wavelengths, refractive_indexes, 
                                                    alpha_chl, alpha_w, beta_w, alpha_wc, solar_energy) 
                                  for i in range(m) 
                                  for j in range(n)]

                    OSA = np.asarray(dask.compute(zr)).reshape((m, n, 2))
                    nlevels=np.arange(0.01,0.04,0.001)

                    irradiance_water = (rads[:,:,0]*OSA[:,:,0]+rads[:,:,1]*OSA[:,:,1])/(OSA[:,:,0]+OSA[:,:,1])

                    print("Time to finish {} with mean OSA {}".format(datetime.datetime.now()-startdate,
                          np.mean(irradiance_water)))

                    # Write to file
                    data_array=xr.DataArray(data=irradiance_water,dims={'lat':lat,'lon':lon})
                    if not os.path.exists("ncfiles"):
                        os.mkdir("ncfiles")
                    data_array.to_netcdf("ncfiles/irradiance.nc")

因为我需要为3条社会经济路径(SSP)的几个CMIP6模型运行这个脚本。对于每个模型和SSP，我必须计算150年的月光值，140个波长的光谱值，在1×1度分辨率的全球范围内。这是CPU和内存消耗，我想知道是否有改进我的矢量化的方法或更好的使用达斯克的方法。如果有人能为我指出如何提高加速比的正确方向，那就太好了。

Answer 1

这些将不会影响业绩，但仍有助于解决以下问题：

类型提示

一些疯狂的猜测，但是：

def calculate_OSA(
    µ_deg: float,
    uv: float,
    chl: float,
    wavelengths: ndarray,
    refractive_indexes: ndarray,
    alpha_chl: float,
    alpha_w: float,
    beta_w: float,
    alpha_wc: float,
    solar_energy: float,
):

尽管如此，考虑到参数的数量很大，使用类型化成员创建一个@dataclass并将其作为参数传递或在其上创建一个方法可能更容易。

No-op返回

您的最后返回可以删除。但令人怀疑的是，alpha_diffuse、alpha_direct_chl和alpha_diffuse_chl都未使用。看一下您的Github，这里似乎忘记了将调用复制到calculate_spectral_and_broadband_OSA。

Answer 2

看看jupyter笔记本，我想知道一点缓存是否有帮助？这些数据点中有多少是真正独特的？像在回忆录装饰器中包装经常调用的函数这样简单的事情可能会有所帮助。任何只使用浮动的calculate_函数都是很好的选择--我不认为回忆录任何需要向量的东西都会有帮助。