python-xarray - 以 block 的形式创建 xarray DataArray 并将其写入 NetCDF

Question

是否也可以创建一个核外 DataArray，并使用 xarray 将其逐块写入 NetCDF4 文件？

例如，当维度更大时，我希望能够以核外方式执行此操作，因此我无法将整个数组存储在内存中:

num_steps = 20
num_times = 100
#Create DataArray
d = xr.DataArray(np.zeros([num_steps, num_times], np.float32),
                 {'Step': np.arange(num_steps),
                  'Time': np.arange(num_times)},
                 ('Step', 'Time'))
#Computatation
for i in range(num_steps):
    d[i, :] = i
#Write to file
d.to_netcdf('test.nc')

所以我不想在内存中创建整个 NumPy 数组，我希望计算和写入文件阶段一次完成一个块(在本例中在 Step 维度上分块)。

更新:
似乎(来自@jhamman 的回答)可能无法使用 xarray 实现我上面的示例。我的主要兴趣是加深对 xarray 的核外计算的理解，所以我没有要问的特定计算，但是，由于有人要求我提供一个更复杂的示例，因此我有一个潜在的应用程序有是:

for i in range(num_steps):
    u[:] = f(u)
    s[:] = g(s)
    d[i, :] = u[:] * s[:]

哪里 u和 s是维度 Time 的 xr.DataArrays，和 f和 g是仅依赖于上一步输入数组的 PDE 求解器。假设有 1000 步，但是时间维度太大，我只能在内存中存储一​​两个，因此赋值给 d必须写入磁盘，然后释放相关的内存。

Answer 1

是的，xarray 支持核外数组和块写入。您将需要使用 xarray 操作和 Dask 编写计算。数组而不是 NumPy 数组。 xarray docs在这里应该会有所帮助。

更新 :对于这样的模拟，您需要计算每个函数 f使用 dask.delayed .然后您可以使用 dask.array.from_delayed 将结果转换为 dask 数组，将它们包装回 xarray.DataArray并使用 to_netcdf() 将数据直接写入磁盘.结果以流方式进行，f()和 g()并行计算并且加载到内存中的时间步骤不超过几个:

import dask
import dask.array as da
import numpy as np
import xarray

def f(x):
    return 1.1 * x

def g(x):
    return 0.9 * x

num_steps = 1000
num_times = int(1e6)

u = np.ones(num_times)
s = np.ones(num_times)

arrays = []
for i in range(num_steps):
    u = dask.delayed(f)(u)
    s = dask.delayed(g)(s)
    product = da.from_delayed(u * s, shape=(num_times,), dtype=float)
    arrays.append(product)

stacked = da.stack(arrays)
data_array = xarray.DataArray(stacked, dims=['step', 'time'])
%time data_array.to_netcdf('results.nc')
# CPU times: user 7.44 s, sys: 13.5 s, total: 20.9 s
# Wall time: 29.4 s

你会注意到 xarray 对这个计算来说非常外围:大部分计算是用 dask/numpy 完成的。您也可以使用 xarray 对象轻松地做到这一点，但我们没有一种方便的方法来通过 dask 延迟对象传递标记的数组元数据，因此无论哪种方式，您都需要在另一侧重建元数据。

您可能会争辩说，在这里使用 dask 是矫枉过正，而且您可能是对的。即使您想使用 dask 进行并行化，您仍然可能希望在每个步骤之后以有效 netCDF 文件的形式检查点模拟。

因此，您可能需要一个在每次迭代时扩展 netCDF 文件的简单循环。这是 not yet supported通过 xarray 但这将是一个很好的功能。像下面这样的界面应该是可能的:

for i in range(num_steps):
    u[:] = f(u)
    s[:] = g(s)
    d[:] = u[:] * s[:]
    d.to_netcdf('results.nc', extend='step')

同时，您可以为每个步骤编写单独的文件，例如，

for i in range(num_steps):
    u[:] = f(u)
    s[:] = g(s)
    d[:] = u[:] * s[:]
    d.to_netcdf('results-%04d.nc' % i)

然后，您可以将所有数据一起加载，然后使用 open_mfdataset 将其合并到一个文件中。 ，例如，

combined = xarray.open_mfdataset('results-*.nc', autoclose=True)
combined.to_netcdf('results-combined.nc')