Python中没有循环的多维Xarray索引答案

【问题标题】：Multidimensional Xarray indexing without loops in PythonPython中没有循环的多维Xarray索引
【发布时间】：2018-04-10 21:41:40
【问题描述】：

我正在查看海洋模型的输出，我想从输出中创建一个水柱底部的温度网格（x,y 相当于 lat,lon），即最深的网格单元.在 xarray 数据集中，我有最大深度（见下文，“深度”）。

我可以用一个非常慢的循环来做到这一点，但想知道是否有办法避免循环，或者至少是部分循环。

到目前为止，代码的数据如下所示：

# load data as xarray
data_dir = 'run04'
ds1 = open_mdsdataset(data_dir,iters=np.arange(0,10001,5000),prefix=['U','V','W','S','T','Eta'])
ds1 = ds1.rename({'T':'Tt'}) # T doesn't work because it thinks its transpose

ds1 的外观：

<xarray.Dataset>
Dimensions:  (XC: 40, XG: 40, YC: 30, YG: 30, Z: 100, Zl: 100, Zp1: 101, Zu: 100, time: 3)
Coordinates:
  * XC       (XC) >f4 2500.0 7500.0 12500.0 17500.0 22500.0 27500.0 32500.0 ...
  * YC       (YC) >f4 2500.0 7500.0 12500.0 17500.0 22500.0 27500.0 32500.0 ...
  * XG       (XG) >f4 0.0 5000.0 10000.0 15000.0 20000.0 25000.0 30000.0 ...
  * YG       (YG) >f4 0.0 5000.0 10000.0 15000.0 20000.0 25000.0 30000.0 ...
  * Z        (Z) >f4 -7.0 -21.0 -35.0 -49.0 -63.0 -77.0 -91.0 -105.0 -119.0 ...
  * Zp1      (Zp1) >f4 0.0 -14.0 -28.0 -42.0 -56.0 -70.0 -84.0 -98.0 -112.0 ...
  * Zu       (Zu) >f4 -14.0 -28.0 -42.0 -56.0 -70.0 -84.0 -98.0 -112.0 ...
  * Zl       (Zl) >f4 0.0 -14.0 -28.0 -42.0 -56.0 -70.0 -84.0 -98.0 -112.0 ...
    rA       (YC, XC) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    dxG      (YG, XC) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    dyG      (YC, XG) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    Depth    (YC, XC) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    rAz      (YG, XG) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    dxC      (YC, XG) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    dyC      (YG, XC) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    rAw      (YC, XG) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    rAs      (YG, XC) >f4 dask.array<shape=(30, 40), chunksize=(30, 40)>
    drC      (Zp1) >f4 dask.array<shape=(101,), chunksize=(101,)>
    drF      (Z) >f4 dask.array<shape=(100,), chunksize=(100,)>
    PHrefC   (Z) >f4 dask.array<shape=(100,), chunksize=(100,)>
    PHrefF   (Zp1) >f4 dask.array<shape=(101,), chunksize=(101,)>
    hFacC    (Z, YC, XC) >f4 dask.array<shape=(100, 30, 40), chunksize=(100, 30, 40)>
    hFacW    (Z, YC, XG) >f4 dask.array<shape=(100, 30, 40), chunksize=(100, 30, 40)>
    hFacS    (Z, YG, XC) >f4 dask.array<shape=(100, 30, 40), chunksize=(100, 30, 40)>
    iter     (time) int64 dask.array<shape=(3,), chunksize=(1,)>
  * time     (time) int64 0 5000 10000
Data variables:
    Eta      (time, YC, XC) float32 dask.array<shape=(3, 30, 40), chunksize=(1, 30, 40)>
    V        (time, Z, YG, XC) float32 dask.array<shape=(3, 100, 30, 40), chunksize=(1, 100, 30, 40)>
    W        (time, Zl, YC, XC) float32 dask.array<shape=(3, 100, 30, 40), chunksize=(1, 100, 30, 40)>
    S        (time, Z, YC, XC) float32 dask.array<shape=(3, 100, 30, 40), chunksize=(1, 100, 30, 40)>
    U        (time, Z, YC, XG) float32 dask.array<shape=(3, 100, 30, 40), chunksize=(1, 100, 30, 40)>
    Tt       (time, Z, YC, XC) float32 dask.array<shape=(3, 100, 30, 40), chunksize=(1, 100, 30, 40)>

以及获取最深单元温度值的循环：

# find the deepest wet cell at each gridpoint
# loop through timesteps 
t_at_bottom1 = np.zeros((ds1.time.size,ds1.YC.size,ds1.XC.size))
for ti in np.arange(0,ds1.time.size,1):
    # loop through x,y indices
    for yi in np.arange(0,ds1.YC.size,1):        
        for xi in np.arange(0,ds1.XC.size,1):
            # look for the grid cell closest to the bottom
            t_at_bottom1[ti,yi,xi] = ds1.Tt.sel(time=ds1.time[ti],Z=-ds1.Depth.values[yi,xi],YC=ds1.YC[yi],XC=ds1.XC[xi],method='nearest')

感谢您的帮助。

【问题讨论】：

Depth 不依赖于时间，因此您至少应该能够摆脱这个循环。此外，如果这是使用 dask 的数据集的实际大小，则可能是矫枉过正。你试过ds.load()吗？
这可能是真的@mathause。如果我只是使用第一个时间步循环遍历 x 和 y，我可以获得每列单元格的底部单元格的 z 坐标中的索引。正如您所说，这对于后续的时间步长不会改变。但是，我不知道如何使用 z 坐标的输出作为后续时间步长的温度 (Tt) 数组中的索引。这是一个粗略的模型运行，因此后续运行将拥有更大的数据集。不过感谢您的提示。

标签： python indexing nested-loops dask python-xarray

【解决方案1】：

看看vectorized indexing。

2D 中的简单示例：

import xarray as xr
import numpy as np
data = np.arange(12).reshape(3, 4)
da = xr.DataArray(data, dims=['depth', 'x'],
                  coords=dict(depth=[0, 1, 2], x=[0, 1, 2, 3]))

da 看起来像

<xarray.DataArray (depth: 3, x: 4)>
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
Coordinates:
  * depth    (depth) int64 0 1 2
  * x        (x) int64 0 1 2 3

现在可以通过索引与另一个数据数组一起选择与 x 相关的深度：

sel = xr.DataArray([0, 1, 0, 2], dims=['x'])
da[sel]

<xarray.DataArray (x: 4)>
array([ 0,  5,  2, 11])
Coordinates:
    depth    (x) int64 0 1 0 2
  * x        (x) int64 0 1 2 3

这需要xarray version 0.10.0 or later。

【讨论】：