Python xarray/pandas 在多个网格点评估函数

Question

我正在尝试在 240x262 网格的所有网格点评估一个函数，但这需要一些优化。我是 Python 新手；通常我只会使用 for 循环遍历所有网格点，并在每个点上逐一评估函数。但可以肯定的是，可以考虑使用 python 更好的方法，我尝试了很多东西但没有成功:-(我希望有人可以帮助我或让我走上正确的轨道。

下面是我会使用的代码示例，但该函数太复杂，无法在合理的时间内完成。所以我有两个问题：

1. 在评估这个数组时，有没有比使用 for 循环更好的方法？
2. 我不需要在所有个网格点评估函数。我可以使用 NaN 轻松掩盖它吗？这样该函数仅用于非 NaN 网格点，而不实际进入该网格点（即使用 for 循环）。 X 和 Y 只是 1D（时间轴）数组，没有什么花哨的，来自 dt x nlat x mlon 3D 数组。也许你也可以生成一个直接插入3D字段的函数......

nlat = 240
mlon = 262
A = np.zeros([nlat,mlon])

for i in np.arange(0,nlat-1):
    for j in np.arange(0,mlon-1):
        A[i,j] = function(x,y)
        
            
df = xr.DataArray(A)
df.to_netcdf('A.nc')

谢谢！

Answer 1

在学习 xarray 之前，你应该从基础开始学习 numpy。

import numpy as np
import xarray as xr

def function(x, y):
    return x - y

x = np.random.rand(240, 1)
y = np.random.rand(262)

df = xr.DataArray(np.vectorize(function)(x, y))

可能还有其他更有效的方法，但您可以在 numpy 中对函数进行矢量化，然后将数组提供给它。由于您希望它们变为 2D，因此您需要将其设为垂直列（如果无法以这种方式创建，可以使用 np.reshape）。

你会得到一个 numpy 数组作为结果，并可以使用它来创建一个 xr DataArray。