|
+ v I0 m7 \, l7 e( H u/ b 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 # [. r1 Y. b& x C! t" f4 y# m
$ r6 U9 w& q+ v2 p4 v 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 0 N! x/ n% O1 p9 c
; g2 L9 e: Z H8 E3 i$ p9 q' A- ` 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 ; B8 O* U" z! e" z
4 c+ E9 B. G9 h, }6 W1 T c 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
9 S1 a( h# `" R+ a3 j" h4 D 其中几个关键点解释一下：
. B; l8 n: Z0 e& H, _ （1）首先定义拟合正太分布的函数 6 ], `' i4 n, u# }- U1 a+ H
def norm_fit(data):
2 ^) U G1 A; d" m' {. ? loc, scale = st.norm.fit(data)+ U% [, C% c0 ^, Z
return np.array([loc, scale]) 0 O8 Y, I/ A1 w' M3 s0 L6 q
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 3 X, p' u6 O" b% @& x
（2）xarray分块 5 ~; S0 L4 T3 L+ \/ \5 e; N
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])1 ~3 E$ e: _+ T3 ~4 h+ u
x = x.chunk({"lev":5, "lat":100,"lon":100})
3 E* }+ S1 e( W4 H9 u xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 0 ^5 {' i/ G: z: q! P1 Y F
（3）xarray.apply_ufunc函数
: m; x# [8 L* w% Y/ @ o result = xr.apply_ufunc(norm_fit,' p3 {# w: v: } }
x,7 X2 a% a) I+ u' w* S
input_core_dims=[["time"]],
8 Y6 O1 n6 q. w( z4 H9 T+ k3 g# ^ output_core_dims=[["paras"]],
" J. G+ g0 X' N dask="parallelized",1 C# `) F% k% {6 y/ `9 r
output_dtypes=[np.float], F& h+ c9 e' Y H5 w+ b y3 B4 c
dask_gufunc_kwargs={output_sizes:{"paras":2}}; l0 |. @4 C3 _% w6 T
) 8 H- A7 d0 u- S
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 6 c" j5 E) j3 o8 e. z& c) {/ i* r
) {. R, J" v+ O/ ]5 ` 以下是全部代码：
( k" t9 v1 Q" b3 v from scipy import stats as st; ?# I% {& |' l, u& W+ O, A& _
0 F4 ^) Z* |) g1 U& @" d! M7 ^
import xarray as xr
8 {& Y4 @ v" V import dask
3 t9 `8 m. z. ?! D# { import numpy as np
9 i1 h6 N6 h. h3 ?' W$ n" k9 N from dask.diagnostics import ProgressBar9 o* g! u4 H0 M* C4 z' i
/ ^/ c. Q0 ]/ |4 | def norm_fit(data):! J/ m4 t9 o0 @0 |& F# c) d2 s1 O
loc, scale = st.norm.fit(data)
: O1 i5 _# A: J+ x return np.array([loc, scale])
9 l9 C# |' M& G% i# \+ T( A
# U8 u n! ?4 z) H% O. P rs = np.random.RandomState(0)6 k' t/ s/ J m+ \
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
8 a* Y9 C7 x1 H/ N% a0 R4 D- w x = x.chunk({"lev":5, "lat":100,"lon":100})+ C+ L* l( a: }" ^, ~; Y8 e$ D
5 P3 T& V* }9 t; _% M #使用apply_ufunc计算，并用dask的并行计算
7 M& |# R q& D5 U result = xr.apply_ufunc(norm_fit,
& } K# M9 D, s% E1 G x,
% Y7 h! P! S- f: J! q input_core_dims=[["time"]],
, k$ b+ S: A) k6 t# k: N1 f! G. S output_core_dims=[["paras"]],* j: w! s) _+ F# n' N z
dask="parallelized",
; U; p" O2 t e' [ output_dtypes=[np.float],
' T! A1 M0 L/ d/ V dask_gufunc_kwargs={output_sizes:{"paras":2}}
! ]( }! x/ c9 g )) Y7 C6 X5 T/ ~" G# w
+ Y \' S5 I. O, _# A #compute进行真实的计算并显示进度) Y; K9 t6 M1 f/ l
with ProgressBar():
% D0 V% g0 g }+ r* e result = results.compute()
" d2 v8 M w( U E7 g( l, O" f( k7 f/ @
#结果冲命名保存到nc文件+ r# f9 d% `+ `( I0 d. ?# Z# K
result = result.rename("norm_paras")
: j' q/ h( a( M result = result.to_netcdf("norm_fit_paras.nc",compute=False)+ c1 `- Y3 g+ J& Y2 F
with ProgressBar():
% {: ~8 U( V' F/ P result.compute() - I! T. l3 k1 H( e
转自：气海同途
! t0 ~" E6 W) Y, @* Y 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
) R( T/ `* P& O* O! y# i $ N; H: Q& x, f- U
推荐：
) w: p* h- G1 a 1、Python语言在地球科学领域中的应用实践技术应用
3 a& E0 f1 A; a: q- _ 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” ) E/ n p( L8 w- b; T1 J
3、Python在气象与海洋中的实践技术应用精品课程
7 s' M, r/ u8 n# e 4、Python在WRF模型自动化运行及前后处理中的实践技术应用
' M) P9 ^! `0 F. C8 d2 d 5、全套Python机器学习核心技术与案例分析实践应用视频 2 U- `) \* H$ ~' o. ^
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
: j" W$ D! P4 a) U 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
4 w0 ~ j- ^" `6 Q* _5 e4 f0 D+ n3 M2 H5 m/ ^* G7 \$ E+ H( `
( ~5 x7 h+ J5 e3 l: `/ M# w
! a: b3 z _ a3 }6 _
8 X& M# Y- V/ _) n: G | 
