xarray+dask处理大规模海洋气象数据

气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。

& X7 ?, L8 H0 h5 L" l, X

大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。

- B" O% k) Q2 G

本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。

" o, W& O ^2 `' M7 w7 h

数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。

其中几个关键点解释一下：

$ }5 A6 ~! ~, V9 y1 l9 q( w

（1）首先定义拟合正太分布的函数

def norm_fit(data): / n" ~0 J V( U1 z ~1 @ loc, scale = st.norm.fit(data) return np.array([loc, scale])

这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。

8 x5 Y+ m! P* F! p$ }( ^0 `- T

（2）xarray分块

^2 ]6 B( `6 ~* C; [

x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) / U& g7 D+ P6 H& h8 @ x = x.chunk({"lev":5, "lat":100,"lon":100})

9 {$ r, g* f& T" }

xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。

（3）xarray.apply_ufunc函数

result = xr.apply_ufunc(norm_fit, 0 L7 L1 N7 K! V& S0 C" g x, 8 \& _; @$ z D1 x; z+ @ input_core_dims=[["time"]], : U- ]& f* T" n& s' N, m output_core_dims=[["paras"]], dask="parallelized", 5 \2 j) ?" n7 d, E output_dtypes=[np.float], * d: {0 v2 U0 V4 q( l, ^- J dask_gufunc_kwargs={output_sizes:{"paras":2}} )

apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）

$ @! X* I/ Q" _9 {) T

以下是全部代码：

from scipy import stats as st % |9 l( f/ F9 _7 g5 i import xarray as xr import dask import numpy as np from dask.diagnostics import ProgressBar + f, K; y2 O& o; B" j def norm_fit(data): loc, scale = st.norm.fit(data) + s( }. Z7 [* H- V( E- B return np.array([loc, scale]) l+ c8 p( G; z8 T+ g% K" L( ^ l ) }+ W) `6 _" E- j/ l rs = np.random.RandomState(0) x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) 0 q1 ?; Y: j9 ^- { x = x.chunk({"lev":5, "lat":100,"lon":100}) #使用apply_ufunc计算，并用dask的并行计算 result = xr.apply_ufunc(norm_fit, x, # j# n$ F1 W* d/ W5 @; [ input_core_dims=[["time"]], output_core_dims=[["paras"]], % M+ B3 g' Q \ dask="parallelized", ! J4 x" k: ]7 T output_dtypes=[np.float], [1 f3 E6 `# T9 S dask_gufunc_kwargs={output_sizes:{"paras":2}} 2 C7 {. }7 @ _2 M& ^* B ) / W# Y4 }$ Z6 K% W3 a' T #compute进行真实的计算并显示进度 with ProgressBar(): result = results.compute() - Q3 p1 o+ I) j #结果冲命名保存到nc文件 " t( H5 M" C1 y% H result = result.rename("norm_paras") result = result.to_netcdf("norm_fit_paras.nc",compute=False) with ProgressBar(): result.compute()

转自：气海同途

关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料