xarray 多维切片数据

Question

我正在使用来自 GFS 的预测数据。我编写了以下函数来从预测档案中提取时间序列：

def time_series_from_ensemble_archive(ensemble_archive, lead_time: int=0, ensemble_member: int=0):

    data = ensemble_archive

    data['age'] = (data.validityDateTime - data.as_of_datetime).astype(np.float)
    age_idx =  data['age']==lead_time

    return data[:, ensemble_member, :, :, :].values[age_idx.T, :, :]

它按预期工作：

这里是 data:

Coordinates:
  * validityDateTime    (validityDateTime) datetime64[ns] 2017-10-01 ...
  * perturbationNumber  (perturbationNumber) int32 0 1 2 3 4 5 6 7 8 9 10 11 ...
  * lon                 (lon) float64 -119.0 -118.5 -118.0 -117.5 -117.0 ...
  * lat                 (lat) float64 45.5 45.0 44.5 44.0 43.5 43.0 42.5 ...
  * as_of_datetime      (as_of_datetime) datetime64[ns] 2017-10-01 ...
Attributes:
    name:                 2 metre temperature

还有我的函数：

temp_ts = time_series_from_ensemble_archive(data)
temp_ts.shape
(124, 10, 20)

type(temp_ts)
numpy.ndarray

但是，我觉得这不是最 'pythonic' 或 'xarrayic' 的方法，return 另一个 xarray 对象会更好。这里有改进建议吗？有人可以使用 expand_dims 或 .sel 方法提供解决方案吗？

Answer 1

xarray提供了多种方式index and select data. You might try indexing with dimension names，例如：

# select using positional & boolean indices
return data[{
    'perturbationNumber': ensemble_member,
    'validityDateTime': (data['age'] == lead_time)}]

或者如果 lead_time 实际上是一个位置索引，就

# select using positional indices
return data[{
    'perturbationNumber': ensemble_member,
    'validityDateTime': lead_time}]

如果您想提供索引标签而不是它们的位置，您可以只使用 .sel 或 .loc 方法：

# select using labels
return data.sel(
    perturbationNumber=ensemble_member,
    validityDateTime=lead_time)

或

# select using labels and boolean indices
return data.loc[{
    'perturbationNumber': ensemble_member,
    'validityDateTime': (data['age'] == lead_time)}]

调用da.values是returnsxarray数据的numpy数组后端的步骤。您的代码没有理由不能使用您提供的索引实际 xarray DataArray（没有 .values）的索引。

xarray 多维切片数据

xarray slice data by multiple dimensions

slice

multidimensional-array

python-xarray