为 xarray 中的每个 lat/lon 点选择给定日期的值

Question

我有一个 xr.DataArray 对象，网格上的每个经纬度点都有 2015 年的一天（作为 cftime.DateTimeNoLeap 对象）。

date_matrix2015

    <xarray.DataArray (lat: 160, lon: 320)>
array([[cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0)],
       ...,
       [cftime.DatetimeNoLeap(2015, 3, 14, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 3, 14, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 3, 14, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0)]], dtype=object)
Coordinates:
    year     int64 2015
  * lat      (lat) float64 -89.14 -88.03 -86.91 -85.79 ... 86.91 88.03 89.14
  * lon      (lon) float64 0.0 1.125 2.25 3.375 4.5 ... 355.5 356.6 357.8 358.9

我在同一经纬度网格上还有另一个 xr.DataArray 用于垂直速度 (omega)，其中包含 2015 年每一天的数据。在每个经纬度点，我想 selectdate_matrix2015对应日期的速度值。理想情况下，我想做这样的事情： omega.sel(time=date_matrix2015)

我曾尝试通过迭代手动构建新的数据数组，但运气不佳。

有人有什么想法吗？提前致谢！

------------编辑----------------

这是该问题的最小可重现示例。为了阐明我在寻找什么：我有两个 DataArrays，一个用于每日降水值，一个用于每日 omega 值。我想为每个 lat/lon 点确定出现最大降水的那一天（我想我已经正确完成了这一部分）。从那里我想在每个 lat/lon 点 select 最大降水日出现的 omega 值。所以最终我想得到一个欧米茄值的 DataArray，它有两个维度，纬度和经度，其中每个 lat/lon 点的值是该位置最大降雨日的欧米茄值。

import numpy as np
import xarray as xr
import pandas as pd

precip = np.abs(8*np.random.randn(10,10,10))
omega = 15*np.random.randn(10,10,10)
lat = np.arange(0,10)
lon = np.arange(0, 10)
##Note: actual data resolution is 160x360
dates = pd.date_range('01-01-2015', '01-10-2015')

precip_da = xr.DataArray(precip).rename({'dim_0':'time', 'dim_1':'lat', 'dim_2':'lon'}).assign_coords({'time':dates, 'lat':lat, 'lon':lon})
omega_da = xr.DataArray(omega).rename({'dim_0':'time', 'dim_1':'lat', 'dim_2':'lon'}).assign_coords({'time':dates, 'lat':lat, 'lon':lon})

#Find Date of maximum precip for each lat lon point and store in an array
maxDateMatrix = precip_da.idxmax(dim='time')


#For each lat lon point, select the value from omega_da on the day of maximum precip (i.e. the date given at that location in the maxDateMatrix)

Answer 1

您可以将 da.sel with da.idxmax 与 select 配对 index 沿任意维度的最大值：

In [10]: omega_da.sel(time=precip_da.idxmax(dim='time'))
Out[10]:
<xarray.DataArray (lat: 10, lon: 10)>
array([[ 17.72211193, -16.20781517,   9.65493368, -28.16691093,
         18.8756182 ,  16.81924325, -20.55251804, -18.36625778,
        -19.57938236, -10.77385357],
       [  3.95402784,  -5.28478105,  -8.6632994 ,   2.46787932,
         20.53981254,  -4.74908659,   9.5274101 ,  -1.08191372,
          9.4637305 , -10.91884369],
       [-31.30033085,   6.6284144 ,   8.15945444,   5.74849304,
         12.49505739,   2.11797825, -18.12861347,   7.27497695,
          5.16197504, -32.99882591],
       ...
       [-34.73945635,  24.40515233,  14.56982584,  12.16550083,
         -8.3558104 , -20.16328749, -33.89051472,  -0.09599935,
          2.65689584,  29.54056082],
       [-18.8660847 ,  -7.58120994,  15.57632568,   4.19142695,
          8.71046261,   9.05684805,   8.48128361,   0.34166869,
          8.41090015,  -2.31386572],
       [ -4.38999926,  17.00411671,  16.66619606,  24.99390669,
        -14.01424591,  19.85606151, -16.87897   ,  12.84205521,
        -16.78824975,  -6.33920671]])
Coordinates:
    time     (lat, lon) datetime64[ns] 2015-01-01 2015-01-01 ... 2015-01-10
  * lat      (lat) int64 0 1 2 3 4 5 6 7 8 9
  * lon      (lon) int64 0 1 2 3 4 5 6 7 8 9

请参阅 xarray docs on Indexing and Selecting Data for more info, especially the section on Advanced Indexing 的重要部分，其中介绍了使用 DataArrays 作为索引器以进行强大的整形操作。

为 xarray 中的每个 lat/lon 点选择给定日期的值

Selecting the value at a given date for each lat/lon point in xarray

indexing

python-xarray