如何从 xarray 中的现有变量计算新变量

Question

我有数据，包括多个受试者随时间推移的几个变量的测量值。它位于 3D x阵列中。DataArray，其中第一维索引主题，第二维索引变量，第三维索引时间。如何有效地添加一个新变量，该变量是每个时间步长当前变量的函数？例如，假设我们有这个数组

import numpy as np
import xarray as xr
x = xr.DataArray(
  np.random.randn(2,3,4), 
  dims=("unit", "variable", "time"), 
  coords={"variable": list("ABC")}
)

如何修改x使其具有形状(2,4,4)，并且新变量由每个主题的每个时间步的变量 A、B、C 的范数组成？

编辑：

我想出一种方法来做到这一点：

norm = np.sqrt((x*x).sum('variable'))
tfmd = norm.assign_coords(variable='norm').expand_dims('variable', 1)
x = xr.concat((x, tfmd), dim='variable')

有没有更通用的方法来完成第一行？有没有办法缩短后两行？

Answer 1

Hei，如果您有 xarray 数据集，则基于旧变量添加另一个变量很简单。特维辛、马库斯·

"""
add new variable based on old ones in xarray
adapted from http://xarray.pydata.org/en/stable/generated/xarray.Dataset.assign.html
"""
import numpy as np
import xarray as xr
ds = xr.Dataset(
    {
        "temperature_1": (
            ("lat", "lon"),
            20 * np.random.rand(4).reshape(2, 2),
        ),
        "temperature_2": (("lat", "lon"), 20 * np.random.rand(4).reshape(2, 2)),
    },
    coords={"lat": [10, 20], "lon": [150, 160]},
)
print(ds)
ds = ds.assign(temperature_3=ds["temperature_1"] - ds["temperature_2"])
print(ds)