xarray 和 dask：高效处理大型 netcdf 文件

Question

我正在尝试对一个非常大的 netcdf 文件进行简单的计算，并且正在努力加快速度 - 可能是因为我主要使用 julia 和 R 进行编程。我认为 xarray/dask 是解决这个问题的最佳方法，但我正在努力让它发挥作用。

netcdf 文件具有三个维度：纬度 (4320)、经度 (8640) 和时间 (792)。它是全球降水数据集。我想计算一组 129,966 个纬度/经度点位置的年降水量平均值（时间维度以月为单位，因此每 12 个）。

这是我的粗略骨架：

import xarray as xr
import netCDF4 as nc
import numpy as np
import pandas as pd
from dask import delayed, compute
from dask.distributed import Client
import dask
from tqdm import tqdm

num_workers = 10  # Adjust this to the number of cores you want to use
threads_per_worker = 1  # Adjust as needed

client = Client(n_workers = num_workers, threads_per_worker = threads_per_worker, dashboard_address=':8787')

## this is the netcdf file
ds = xr.open_mfdataset(f'../data/terraclim/TerraClimate_{var}_*.nc', chunks='auto')

lon = ds['lon'].values
lat = ds['lat'].values

# Initialize arrays for indices
lonindx = np.empty(coords.shape[0], dtype=int)
latindx = np.empty(coords.shape[0], dtype=int)

# Loop through each row of coords, a pandas dataframe of lat and lon coordinates to extract data from
for i in tqdm(range(coords.shape[0]), desc="Processing"):
    lonindx[i] = np.argmin(np.abs(lon - coords.at[i, 'LON']))  
    latindx[i] = np.argmin(np.abs(lat - coords.at[i, 'LAT']))  

## indexes the netcdf file (ds) by lat and lon and returns an array of year-means
def compute_means(i):

    vals = ds[var][0:791, latindx[i], lonindx[i]].values
    means = [np.mean(year) for year in np.array_split(vals, 66)]

    return means

results = []
for i in range(coords.shape[0]):  
    res = dask.delayed(compute_means)(i)
    results.append(res)

results = dask.compute(*results)

但是，当我运行上面的脚本时，它由于内存限制而崩溃。我想我可能缺少一种更合适的方式来使用 xarray 与 dask 的内部兼容性，但我完全迷失了。任何帮助将不胜感激！就像 Julia 或 R 中的解决方案一样。

Answer 1

欢迎来到 Stack Overflow！

您需要处理两个问题，但两者都与您用来处理数据的工具无关。

1. 内存耗尽 典型的年度 TerraClimate 数据变量具有

   8,640 x 4,320 x 12 x 8 = 3,583,180,800 bytes

   。无论您投入多少个 CPU 核心，这都很可能超过除最强大系统之外的任何系统上的可用 RAM。然而，TerraClimate 每年都会生成数据，这些数据正是您在查看

   open_mfdataset()

   命令时所拥有的文件。那么为什么不逐年处理这些文件呢？在我的家用系统（MacBook Pro M2，16GB RAM）上，使用 R 处理单个 var/year 组合（= 1 TerraClimate 文件）需要 43 秒，这意味着您可以在不到一个小时的时间内处理所有 58 年。用

   abind()

   将生成的数组缝合在一起，就完成了。请参阅下面的 R 代码。
2. 数据布局 通常，netCDF 文件具有内部数据布局（简单地说）从最右侧的维度（在本例中为

   lon

   ）到最左侧的维度（

   time

   ）。然后，像在 compute_means() 中所做的那样，在其他维度上分析最左边的维度是非常低效的，因为从文件中读取所需的数据意味着必须遍历整个文件才能获取 8- 中的数据字节块散落各处。当您的 I/O 达到极限时，这就是性能杀手。 （您尝试每年将流程分解为单独的线程，这可能会缓解这一问题，但是嘿！如果各个文件已经按照这种方式组织，为什么首先要采用 MFDataset 方法呢？）
R 中按年份处理

# install.packages("ncdfCF")
library(ncdfCF)
library(abind)

var <- "ppt"
years <- 1958:2024

# This loops over the years and calculates the annual sum of precipitation
# for every lat/lon cell. You could also use `mean` for other variables like
# tmin, tmax.
# The result is a list of matrices, one for each year.
yrs <- lapply(years, function (y) {
  fn <- paste0("../data/terraclim/TerraClimate_", var, "_", y, ".nc"))
  ds <- open_ncdf(fn)
  data <- ds[[var]][]
  apply(data, 1:2, sum, na.rm = TRUE)
})

# Bind the list of matrices into one array, with the third dimension representing years.
all_years <- abind(yrs, along = 3)

你说你想计算“年平均降水量”。以上是年降水量

sum

。如果您想说“年降水量的平均值”，那么您应该像这样取平均值：ppt_multiannual_mean <- apply(all_years, 1:2, mean, na.rm = TRUE)。

请注意，

na.rm = TRUE

至关重要：毕竟，地球表面的 70% 是海洋。