如何使用 DataArray where() 函数根据条件从另一个 DataArray 赋值

Question

我正在使用 xarray 根据来自另一个数据集的值的条件创建一个新的数据集。

输入的Dataset对象ds_season是按季节划分的，具有如下三个维度。

    <xarray.Dataset>
    Dimensions:               (latitude: 106, longitude: 193, season: 4)
    Coordinates:
      * latitude              (latitude) float32 -39.2 -39.149525 ... -33.9
      * longitude             (longitude) float32 140.8 140.84792 ... 150.0
      * season                (season) object 'DJF' 'JJA' 'MAM' 'SON'
    Data variables:
        FFDI 95TH PERCENTILE  (season, latitude, longitude) float64 dask.array<shape=(4, 106, 193), chunksize=(4, 106, 193)>

我需要创建一个具有纬度、经度和时间三个维度的新数据集。经纬度应与输入Dataset的坐标相同，时间坐标应为10年以上的天数。

例如，生成的数据集是这样的：

<xarray.Dataset>
Dimensions:    (latitude: 106, longitude: 193, time: 3653)
Coordinates:
  * latitude   (latitude) float32 -39.2 -39.149525 ... -33.950478 -33.9
  * longitude  (longitude) float32 140.8 140.84792 140.89584 ... 149.95209 150.0
  * time       (time) datetime64[ns] 1972-01-01T00:00:00 1972-01-02T00:00:00 1972-01-03T00:00:00 ... 1981-12-30T00:00:00 1981-12-31T00:00:00
Data variables:
    FFDI 95TH PERCENTILE  (time, latitude, longitude) float64 dask.array<shape=(3653, 106, 193), chunksize=(3653, 106, 193)>

一天的变量应该与当天所在季节的变量相同。这意味着，1972-01-01、1972-02-02 和 1972-02-28 应该具有相同的值DJF 的赛季；并且 1972-04-01、1972-05-02 和 1972-05-31 应该具有与 MAM 季节相同的值。

我正在考虑数据集的 where() 函数，但不知道从哪里开始。 http://xarray.pydata.org/en/stable/generated/xarray.Dataset.where.html?highlight=where#xarray.Dataset.where

Answer 1

首先，一个注释。创建一个新的 DataArray 并在 3 个月内每天复制相同的空间数据会占用大量磁盘空间而没有多大意义。每当您需要特定日期的数据时，我宁愿查询季节 DataArray。 但是，如果您确实需要执行此操作，并且要回答您的问题，我认为最直接的方法是：

1. 首先，新建一个容器； np.ndarray 是一个好主意。
2. 然后，构建日期索引，
3. 查询您的原始季节 DataArray，
4. 最后，创建一个具有时间维度的新 DataArray。

在以下示例中，我创建了一个季节 DataArray 用于测试。如果我完全理解您的问题，您应该能够在第二部分（通过创建 foo）中使用原始数组而无需进行太多更改。

让我们开始吧。首先是导入：

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

创建一个所需大小的空容器。

data_s = np.zeros((4, 10, 10))

用虚拟值填充它。

data_s[0] = 0.5
data_s[1] = 0.9
data_s[2] = 0.8
data_s[3] = 0.45

创建虚拟坐标。

x = y = np.arange(10)

创建季节索引。

seasons = ["spring", "summer", "autumn", "winter"]

最后，创建 DataArray。

bar = xr.DataArray(data_s, coords=[seasons, x, y], dims=['season', 'x', 'y'])

bar 是您要从中提取季节性值的 DataArray。 现在对单个日期重复相同的操作。

创建一个 2000 天的容器数组，我们将填充每个季节的数据。

data = np.ones((2000, 10, 10))
x = y = np.arange(10)
dates = pd.date_range('2000-01-01', periods=2000)

在这里，我假设北方季节从月初开始（借用自 here. 当然，您可以轻松编写更好的函数，例如使用一年中的某一天来获取季节。

season = np.array((dates.month %12 + 3)//3)

创建一个字典，将上面的数字转换为之前在 bar 中分配的季节字符串

seas_to_num = {1:"spring", 2:"summer", 3:"autumn", 4:"winter"}

我们用在 bar[season] 上找到的值填充每一天的数组。

for date, seas in enumerate(season):
    data[date] = bar.sel(season=seas_to_num[seas])

最后，我们创建 DataArray。

foo = xr.DataArray(data, coords=[dates, x, y], dims=['time', 'x', 'y'])

现在选择 4 月 5 日，我们得到了 spring 的值。

In [1]: foo.sel(time=pd.to_datetime("5/4/2001"))
Out[1]: 
array([[0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
   [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9]])
Coordinates:
time     datetime64[ns] 2001-05-03
  * x        (x) int32 0 1 2 3 4 5 6 7 8 9
  * y        (y) int32 0 1 2 3 4 5 6 7 8 9

Answer 2

我同意 Andrea 的观点，即创建一个包含 3653 个不同天且仅复制 4 个不同季节值的数据集通常效率低下。如果您提供更多关于您这样做的更广泛目标的信息，也许我们可以建议一个替代解决方案。

假设你真的想这样做，最快的方法可能是使用 xarray 的groupby broadcasting arithmetic。在下文中，我将假设ds 是您原始帖子中第二个数据集 的名称（维度为(latitude: 106, longitude: 193, time: 3653) 的数据集）。然后你可以很快地做到这一点

zeros = xr.zeros_like(ds)
filled_in = zeros.groupby('time.season') + ds_season

这个建议的灵感来自于我们通常从季节性气候学中计算异常的方式：

# original dataset with dimensions 'time'
ds = xr.open_dataset(...)
# climatology has dimension 'season'
ds_climatology = ds.groubpy('time.season').mean(dim='time') 
# anomaly has dimension 'time'
ds_anomaly = ds.groubpy('time.season') - ds_climatology