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

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

2 Y' Y1 P9 p# i4 I0 T" |/ C4 ^5 ^

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

7 s8 c# L" m1 `- P7 ^" K! X5 a0 X

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

6 Y. N6 b# J8 n

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

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

! ^3 Z( u" ~. v- I

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

def norm_fit(data): ' n0 j; |; e! c9 n& c loc, scale = st.norm.fit(data) 2 r7 j2 n0 D& E6 D return np.array([loc, scale])

0 B3 i* ] g p# t+ F' n8 T% `# X

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

' @6 p$ w2 ^ u6 e9 }& [

（2）xarray分块

, N9 M) [& m+ Z% y- J: s

x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) 7 u) K' a8 l6 q4 a) w3 Z! ] x = x.chunk({"lev":5, "lat":100,"lon":100})

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

5 l( W9 Y$ L2 T+ \* L# O

（3）xarray.apply_ufunc函数

& F$ D$ B4 r* v) ^# ^, [9 V- v3 T! T

result = xr.apply_ufunc(norm_fit, 4 X) q" n# T9 Z; O. a P- R+ M! }; C5 Y x, input_core_dims=[["time"]], . W- ^/ U0 I; m% f2 e- o+ Q output_core_dims=[["paras"]], + T) M/ _/ O. }/ Q dask="parallelized", output_dtypes=[np.float], dask_gufunc_kwargs={output_sizes:{"paras":2}} # q4 K9 g6 D0 a9 n, ]$ C1 H )

* y/ K, T$ G) c" _: ]6 H

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

9 H; K# w) Y8 }

以下是全部代码：

4 h, M4 A5 A( H. a. v/ R

from scipy import stats as st , \6 W# H' n8 u' e8 t5 I9 X import xarray as xr ) j2 x7 e' F+ B9 V' E import dask import numpy as np from dask.diagnostics import ProgressBar 3 a3 _% J% r1 B( m& v def norm_fit(data): 8 f+ d6 z `; l) g2 l! _+ w4 x- g loc, scale = st.norm.fit(data) return np.array([loc, scale]) rs = np.random.RandomState(0) ; _4 K# Z# ?- z9 M& x7 T) V8 e x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) & j: f( W- D- ? ~" v x = x.chunk({"lev":5, "lat":100,"lon":100}) ' J# W% V" C( A; C* Z1 U b) i #使用apply_ufunc计算，并用dask的并行计算 result = xr.apply_ufunc(norm_fit, N( b4 w6 E n% M x, ( N! E1 [; [! l5 X. a: j input_core_dims=[["time"]], output_core_dims=[["paras"]], dask="parallelized", 5 d% e) e* _/ j5 C8 E/ ~ l output_dtypes=[np.float], 2 q9 K) s' v9 `9 T6 j9 C dask_gufunc_kwargs={output_sizes:{"paras":2}} ) #compute进行真实的计算并显示进度 with ProgressBar(): result = results.compute() W0 `0 c& W7 s #结果冲命名保存到nc文件 result = result.rename("norm_paras") result = result.to_netcdf("norm_fit_paras.nc",compute=False) with ProgressBar(): result.compute()

转自：气海同途

9 b, B0 o5 d2 ^9 q7 g

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