在函数中编辑 (xarray) 对象

Question

我正在尝试使用一个函数将额外变量分配给现有数据集“ds”，该函数将数据集作为参数并且应该 return 调整后的数据集：

def Assign_Variable(dataset):
    dataset = dataset.assign(new_var = dataset.x + dataset.y)
    #(and some more manipulations)
    return dataset

然后使用这个函数循环遍历一些数据集来操作它们：

for dataset in [ds]:
    dataset = Assign_Variable(dataset)

然而，如果我现在检查我的数据集 ds，该函数没有做任何事情。

如何在函数中调整我的数据集，return 它们？

Answer 1

我认为你正在寻找的东西可以通过像字典一样处理数据集来完成（例如数据集[new_var] = ...）

作为示例，我有一些代码根据两个风分量 u 和 v 计算风速，并将输出风速 (ws) 保存在现有数据集 (ds) 中：

import xarray as xr
import numpy as np

def windspeed(dataset):
    dataset['ws'] = np.sqrt(dataset['u']**2 + dataset['v']**2)
    return dataset


# Create sample data
lon = np.arange(129.4, 153.75+0.05, 0.25)
lat = np.arange(-43.75, -10.1+0.05, 0.25)
data = 10 * np.random.rand(len(lat), len(lon))
ds   = xr.Dataset({"u": (["lat", "lon"], data), "v": (["lat", "lon"], data)}, coords={"lon": lon,"lat": lat})

# Calc wind speed
ds = windspeed(ds)

# Print output
print(ds)

输出数据集现在包含三个变量 u、v 和 ws。

在函数中编辑 (xarray) 对象

editing (xarray) object in a function

function

python-xarray