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

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

/ Q* s- ]- W) h' q& ]. j

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

6 c3 {; c/ E5 b6 p , q4 @* m" ?8 J# N4 h/ b, A

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

4 x8 X1 v+ a b$ d* g

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

9 t* u Z% w7 `, ^0 O1 B6 h' P$ o

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

8 p& S1 n2 U% I) m: X: z

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

def norm_fit(data): $ |- m0 ] o9 t8 H- L loc, scale = st.norm.fit(data) return np.array([loc, scale])

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

（2）xarray分块

x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) : G; _8 t: z6 q( X9 K) G6 f x = x.chunk({"lev":5, "lat":100,"lon":100})

3 b R9 p3 o$ X

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

, W% g' d8 c- _" }

（3）xarray.apply_ufunc函数

, s# M; ~( P' o9 [1 V* G* E/ @

result = xr.apply_ufunc(norm_fit, % a' l, E& L" U2 a x, input_core_dims=[["time"]], 5 c/ x% X* j+ m8 r& e" g. x output_core_dims=[["paras"]], 9 n! \5 W+ t# ]7 a dask="parallelized", output_dtypes=[np.float], ; W2 A. t$ o6 R: c# q 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: 这里感谢深雨露大佬的指点）

0 A* A1 Q5 q' o

以下是全部代码：

0 z* C, a# W$ G4 p' V8 B8 q! {1 {3 d1 H

from scipy import stats as st import xarray as xr ; g0 j7 j# z# k4 f import dask import numpy as np from dask.diagnostics import ProgressBar def norm_fit(data): loc, scale = st.norm.fit(data) return np.array([loc, scale]) rs = np.random.RandomState(0) - G, S; @: Z: t+ |! A2 W x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) 0 ?4 ?; h5 u' k1 |* r r x = x.chunk({"lev":5, "lat":100,"lon":100}) , _4 X" D4 L1 ~7 Q" Z0 v #使用apply_ufunc计算，并用dask的并行计算 $ V* Q6 K4 P3 w5 e4 L result = xr.apply_ufunc(norm_fit, x, input_core_dims=[["time"]], output_core_dims=[["paras"]], dask="parallelized", output_dtypes=[np.float], dask_gufunc_kwargs={output_sizes:{"paras":2}} # _- l; B! b h6 O ) . A- R$ Z: W( ^; B# D5 ` #compute进行真实的计算并显示进度 # B! k3 V% k: \ with ProgressBar(): 5 I" j7 v' r8 }3 Q9 E5 w! ?/ i# W result = results.compute() " s; W( M: l! D" Y$ r1 a& b; ^6 N 4 A, L6 W) p8 O2 x! Z1 c #结果冲命名保存到nc文件 ( S; S. V) g, t @" U2 X( E& H result = result.rename("norm_paras") result = result.to_netcdf("norm_fit_paras.nc",compute=False) with ProgressBar(): 4 L% Q( s" _" B5 ?8 k/ {9 }9 ] result.compute()

9 H% t/ ~3 ]$ a6 b3 B5 U

转自：气海同途

0 c g% m- T6 Z! O) Z

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