从 dask 块中按名称选择维度

Question

我有一些 grib 格式的集成文件，我想使用 dask 和 xarray 在 Python 中延迟加载。基于https://climate-cms.org/2018/09/14/dask-era-interim.html，我设法按预期延迟加载文件，但现在我想切片并选择维度来绘制一段时间和级别的数据。

更新：我最近回到了这个问题，我终于发现我应该使用da.concatenate，而不是使用da.stack。这个简单的改变解决了我的问题。此问题已相应更新，以防万一有人需要有关如何使用 python（使用 dask 数组用于延迟加载）创建 grib 文件集合的示例，以与使用 GrADS 等软件相同的方式加载和绘制数据。

我的程序如下：

import dask
import dask.array as da
import xarray as xr
import pandas as pd
import numpy as np
from glob import glob
from datetime import date, datetime, timedelta

import matplotlib.pyplot as plt

bpath = '/some/path/to/my/data'

# pressure levels
levels =['1000', '925', '850', '700', '500', '300', '250', '200', '50']

# ensemble member names
ensm = ['M01', 'M02', 'M03', 'M04', 'M05'] 

@dask.delayed
def open_file_delayed(file, vname):
    ds = xr.open_dataset(file, decode_cf='False', engine='pynio')
    return ds

def open_file(file, vname, nlevs, nlats, nlons, ftype):
    file_data = open_file_delayed(file, vname)[vname].data
    return da.from_delayed(file_data, (nlevs, nlats, nlons), ftype)

    # list of files to open (sorted by date)
    # filename mask: PREFIXMEMYYYYiMMiDDiHHiYYYYfMMfDDfHHf.grb
    # MEM: member name (see the levels list)
    # YYYYiMMiDDiHHi: analysis date (passed as an argument to the open_file function)
    # YYYYfMMfDDfHHf: forecast date
    files = sorted(glob(bpath + '/%(dateanl)s/%(mem)s/PREFIX%(mem)s%(dateanl)s*.grb'%
                        {'dateanl': date, 'mem': member}))
   
    ntime = len(files)

    # open the first file in the list to get dimensions and coordinates
    ds0 = xr.open_dataset(files[0], decode_cf='False', engine='pynio')
   
    var0 = ds0[vname]
   
    levs = ds0.lv_ISBL2.data
    lats = ds0.g4_lat_0.data
    lons = ds0.g4_lon_1.data
    
    nlevs = ds0.lv_ISBL2.size
    nlats = ds0.g4_lat_0.size
    nlons = ds0.g4_lon_1.size
    ftype = var0.dtype
   
    ds0.close()
    
    # calculate the date range of the forecasts, in my case len(date_fmt) = 61 (every grib file has 61 times and 9 levels)
    date_fmt = pd.date_range(start=datetime.strptime(date, "%Y%m%d%H"), freq="6H", periods=ntime)
        
    # call the function 'open_file' for all files contained in the list 'files' and stack them up
    dask_var = da.stack([open_file(file, vname, nlevs, nlats, nlons, ftype) for file in files], axis=0)
        
    # xda is the data array with all files
    xda = xr.DataArray(dask_var, dims=['tlev', 'lat', 'lon']) 
            
    # set coordinates values
    xda.coords['time'] = ('time', date_fmt)
    xda.coords['lat'] = ('lat', lats)
    xda.coords['lon'] = ('lon', lons)
    
    return xda

要使用此代码，我会使用（对于单个分析日期 - 202005300 - 2020 年 5 月 30 日，以及一个名为 ZGEO 的变量）：

注意：这部分非常快（需要几毫秒），因为我们只是为实际数据创建一个映射结构，类似于 GrADS 控制文件。

lens_zgeo = [open_ensemble('2020053000', ens, 'ZGEO') for ens in ensm]
dens_zgeo = xr.concat(lens_zgeo, dim='ens')
dens_zgeo.coords['ens'] = ('ens', ensm)

dens_zgeo是一个数据数组，结构如下：

data array structure

从这一点开始，我可以对数据数组的维度进行切片并绘制（这是我最初的意图）：

注意：这部分耗时较长，因为需要从磁盘中读取数据。

dens_zgeo.isel(ens=0,time=0,lev=0).plot()

BOOM，案件结案。谢谢！

Answer 1

我已经用我需要的修改编辑了问题，以获得我想要的结果。对于这种情况，重点是使用da.stack 而不是da.concatenate。通过这样做，我将得到的数据数组连接到我需要的整体维度中。