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数组，我希望Computation和Write to files阶段一次完成一个块（在本例中以Step维度分块）

更新：似乎（来自@ jhamman的回答）可能无法使用xarray实现上面的示例。我主要感兴趣的是用xarray更好地理解核外计算，所以我没有具体的计算，但是，因为我被要求一个更复杂的例子，一个潜在的应用我有：是：

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

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

Answer 1

Dask数组目前不支持项目分配，请参阅Item assignment to Python dask array objects。

如果d是带有dask.array的xarray.DataArray，那么这将不起作用。

~~此外，当前的Xarray后端都没有支持分块写入。~~编辑：正如@shoyer指出的那样，xarray可以逐步写入分块数组。但是，对于您的用例，由于您似乎需要项目分配，可能需要直接使用netCDF4-python库：

from netCDF4 import Dataset

f = Dataset('test.nc', mode='w')
f.createDimension("Step", nsteps)
f.createDimension("time", ntimes)
d = f.createVariable("d", "f4",("Step", "time"))

#Computatation
for i in range(num_steps):
    d[i, :] = i

我认为你的计算比你的例子更复杂，所以你可能会考虑用使用xarray / dask的东西替换= i。

Answer 2

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

更新：对于这样的模拟，您需要使用dask.delayed计算每个函数f。然后，您可以使用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文件的简单循环可能是您想要的。这是xarray的not yet supported，但这将是一个很好的功能。应该可以使用以下界面：

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')

以块

2 个答案: