大NetCDF数据集内存不足与递归错误问题求助

问题背景 #

现有一大型NetCDF数据集，包含192个纬度、382个经度、60265个时间步，数据类型为float64。需要复制该数据集结构以存储复杂计算的结果，但全量加载至内存时会触发内存不足错误。

初始尝试与报错 #

1. 内存不足错误：使用data.copy(deep=True).load()时，触发MemoryError: Unable to allocate 33.1 GiB for an array with shape (192, 384, 60265) and data type float64，无法将33GiB的数据全量加载到内存。
2. 懒加载赋值错误：跳过全量加载直接赋值时，触发NotImplementedError: Lazy item assignment with the vectorized indexer is not yet implemented. Load your data first by .load() or compute().

分块尝试后的新问题 #

采用Dask按time维度分块（每个时间序列为一个块）后，循环赋值操作可以执行，但在加载或保存数据集时触发RecursionError: maximum recursion depth exceeded while calling a Python object，该错误与遍历所有经纬点的循环迭代次数直接相关。

可行解决方案 #

方案1：使用map_blocks结合向量化计算 #

利用xarray的map_blocks和apply_ufunc将逐格点的计算逻辑并行化处理，既避免全量加载内存，又解决循环赋值导致的递归问题。

import xarray as xr
import numpy as np
import sys

def process_grid_cell(time_series):
    # 替换为你的复杂计算逻辑，输入为单格点的时间序列
    # 示例：返回与输入长度一致的全0数组
    return np.zeros_like(time_series)

def doit(data: xr.DataArray) -> xr.DataArray:
    # 调整维度顺序为(lat, lon, time)
    data_transposed = data.transpose("lat", "lon", "time")
    # 对每个lat-lon格点的时间序列应用计算
    result = data_transposed.map_blocks(
        lambda x: xr.apply_ufunc(
            process_grid_cell,
            x,
            input_core_dims=[["time"]],
            output_core_dims=[["time"]],
            vectorize=True,
            dask="parallelized",
            output_dtypes=[x.dtype]
        ),
        template=data_transposed
    )
    return result

def main() -> None:
    # 打开数据集并设置分块，可根据内存调整time分块大小
    ds = xr.open_dataset("dataset.nc", chunks={"time": 60265})['variable']
    result = doit(data=ds)
    # 直接保存，Dask会自动分块计算并写入文件
    result.to_dataset(name='result').to_netcdf('result.nc', mode='w')

if __name__ == "__main__":
    main()
    sys.exit(0)

方案2：预创建空文件，逐格点写入 #

提前创建与原数据集结构一致的空NetCDF文件，循环遍历每个经纬点，计算后直接写入对应位置，全程无需将整个数据集加载到内存。

import xarray as xr
import numpy as np
import sys

def doit(data: xr.DataArray) -> None:
    # 调整维度顺序
    data_transposed = data.transpose("lat", "lon", "time")
    # 创建结构一致的空数据集
    result_ds = xr.Dataset(
        {'result': data_transposed},
        coords=data_transposed.coords
    ).astype(data.dtype)
    # 写入空文件
    result_ds.to_netcdf('result.nc', mode='w')
    
    # 以追加模式打开文件，逐格点写入计算结果
    with xr.open_dataset('result.nc', mode='a') as ds_writer:
        len_lat, len_lon = len(data.lat), len(data.lon)
        for lat_idx in range(len_lat):
            for lon_idx in range(len_lon):
                # 替换为该格点的实际计算逻辑
                grid_result = np.zeros(len(data.time))
                
                # 写入对应位置并保存
                ds_writer['result'][lat_idx, lon_idx, :] = grid_result
                ds_writer.to_netcdf('result.nc', mode='a')

def main() -> None:
    # 懒加载原数据集，无需加载到内存
    ds = xr.open_dataset("dataset.nc")["variable"]
    doit(data=ds)

if __name__ == "__main__":
    main()
    sys.exit(0)

方案3：临时调整递归深度（应急方案） #

如果必须保留原循环赋值逻辑，可以临时增大Python的递归深度，但该方法可能引发稳定性问题，仅作为临时 workaround：

import sys
# 调整递归深度为更大的值，例如33554432
sys.setrecursionlimit(1 << 25)

内容的提问来源于stack exchange，提问作者Benjamin T. Schwertfeger