如何使用 xarray 中的维度和数据来创建新的数据数组？答案

【问题标题】：How can I use dimensions and data in xarray to create a new array of data?如何使用 xarray 中的维度和数据来创建新的数据数组？
【发布时间】：2021-09-28 08:50:17
【问题描述】：

我在Xarray 中遇到了数据集问题。我想对数据集中已有的数据和坐标应用一个函数来计算一组新的数据（本质上，w = f(x, y , z)，其中 x、y 和 z 是浮点数）。

有一系列函数需要应用，并且拉入新变量时需要扩展结果数组的维度。所以f(x)应该是一维数组，g(f(x) , y) 应该是 2D，h(g(f(x), y), z) 应该是从前一个 2D 数组创建的 3D 数组。我可以使用np.outer 来强制进行 1D -> 2D 转换，但在 2D -> 3D 转换时遇到问题。考虑到我遇到的麻烦，我怀疑我是从错误的方向来解决问题的。

下面是一些代码的sn-p。有人能帮我理解从 NumPy/Xarray 的角度解决这个问题的正确方法吗（我真的试图避免陷入循环来完成这项工作……）

import numpy as np
import xarray as xr

data = xr.Dataset(
    data_vars={
        "abs": (("x"), np.ones((200)))
    },
    coords={
        "x": np.linspace(1.5,1.55,200),
        "y": np.arange(5e5),
        "z": np.arange(125)
    }
)

def intensity_decay(da):
    "Return decay curve (exponential decay)."
    init_power = 1e5
    decay = init_power * np.exp(
        -1 * np.outer(da.abs, da.y * 10 ** (-7))
        )
    decay[decay < 1e-3] = 0
    return decay

data["depth_decay"] = (["x", "y"], intensity_decay(data))

def radial_decay(da):
    "Return radial decay curve."
    return np.outer(da.depth_decay, np.exp(-da.z))

radial_decay(data.isel(x=[1, 4, 10])).shape

如您所见，radial_decay 函数不会广播到 z 维度。似乎 Xarray 应该支持这种类型的操作，但我不知道如何解决这个问题，甚至不知道从文档中的哪里开始。我目前有两个要应用这些功能的数据集，所以我可以手动完成，但我想制作一个框架，我也可以在未来的数据上使用。

【问题讨论】：

xarray 可以自动广播命名维度，所以我认为这应该不会太复杂。但据我所知，您对尺寸的初始分配目前有问题。您的示例引发错误，因为您将二维数组作为 abs 的数据参数传递，但仅指定了一个维度 (x)。此外，x 维度的长度在 abs 和 x 坐标之间不匹配。那么absvariable 的维度是多少？
有关在 xarray 中广播的更多信息和示例，请参阅 the documentation 或 this tutorial。
你说得对，我抓了一个样本数组，忘记更改数据。 abs 应该是一维值数组，请参阅编辑。

标签： python numpy python-xarray

【解决方案1】：

玩了一会儿并使用this answer 中的函数结构后，我能够想出正确形状的输出。

通过如下定义radial_decay函数，我可以得到一个具有三个维度的输出数组。

def radial_decay(da):
    "Return radial decay curve."
    da["radial_decay"] = (["x", "y", "z"], da.depth_decay * np.exp(-da.z))
    return da

使用这种结构，radial_decay(data.isel(x=[1, 4, 10])).radial_decay.shape 的输出是(3, 500000, 125)，而不是使用原问题中定义的函数的(1500000,125)。

虽然这似乎可行（我仍然需要验证值是否正确），但我想知道是否还有其他（更好的？）方法可以实现这一目标，特别是如果 just NumPy 有更好的方法来做到这一点（以防我将来由于某种原因无法访问 Xarray）。

【讨论】：