|
! E, o9 z( l! y, P6 i4 `
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 1 j! }. k9 ]5 h, n" Z! |. G9 R
n7 C) a1 p5 v5 Z7 L& _; K
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 6 Z, Z9 I) |" M+ }; s1 I) {% E
2 c: |! P" T% e1 [. W- p
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
$ l2 D7 n8 I, N, ?3 s; k7 g
+ v4 T- o" p* g* D! X* C 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 5 R- s( r! X4 }- q
其中几个关键点解释一下：
$ G$ W- p; ?. x; A （1）首先定义拟合正太分布的函数
. }: x' y) c6 T4 I+ R! m def norm_fit(data):# m( j8 j5 z. r1 p9 l
loc, scale = st.norm.fit(data)
4 J& t2 r) ^% @- }0 U8 _ return np.array([loc, scale]) ) ~6 ] Q" A$ A1 m7 Z/ q, Y* s
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 3 P. S( c2 G# {8 Q* j
（2）xarray分块 6 ^* i4 a K% m4 d
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])# H! s4 I6 l! G4 ~8 m
x = x.chunk({"lev":5, "lat":100,"lon":100})
* ^) u' w, j( r( a8 X/ C xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 2 k* f6 c x# y
（3）xarray.apply_ufunc函数 + q. c1 \8 |9 M1 \- G5 `
result = xr.apply_ufunc(norm_fit,
2 e) t$ M& [ ?" X& B2 p9 ~ x,+ v( ?5 {+ F, ]8 g
input_core_dims=[["time"]],
, m5 O( ~( j; F9 B' l: \ output_core_dims=[["paras"]],& [& @% N# h! V8 r0 y6 D
dask="parallelized",
: s. J1 S9 b6 K: S( q output_dtypes=[np.float],
) g; V. M& e6 u% P, a$ c- Y dask_gufunc_kwargs={output_sizes:{"paras":2}}- W2 h/ j. d1 H1 W, Q0 B
) # H/ G4 ?( ]7 Z; F1 f: Y
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
! O! P+ e9 r- W6 e8 j( d
* m# r0 Y' M1 Y3 Y1 N 以下是全部代码： 1 k" K G0 j6 p) n ^4 t( N0 l3 X
from scipy import stats as st
) n; F. H. |) I% ^6 b3 i
5 Y' B* W% V+ P$ r) Q3 O import xarray as xr, e& L5 \! v% G: H6 Y- ~" a
import dask; q! P& Y' ]: Q% G0 `$ X
import numpy as np, ^4 X5 G2 N; q' ^3 L3 l7 {
from dask.diagnostics import ProgressBar& k2 T, a) c4 W# t8 F
1 u6 Z/ e0 f; z T- ?+ f/ S0 e
def norm_fit(data):& t# [0 Y: J$ a
loc, scale = st.norm.fit(data)
3 v8 \! z: e* T% [7 W; A+ Y- P return np.array([loc, scale])5 n# r+ W; `$ _- G7 J
7 N Y, ^, y* ?, j' L rs = np.random.RandomState(0)
2 J1 R) f: Z, k x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
3 {$ k9 [5 r! O7 e* P" _9 [" C x = x.chunk({"lev":5, "lat":100,"lon":100})- a' A7 D) B4 r+ n1 ?- u
: c7 }: _$ v+ y; X
#使用apply_ufunc计算，并用dask的并行计算% `3 w! {* A% A/ g, N6 \; s. v
result = xr.apply_ufunc(norm_fit,' Q# B- E7 w5 l8 O% l- N- D
x,' z- M! h$ b4 ?& u
input_core_dims=[["time"]],
& L# a `9 B$ d7 q [% u' U* m output_core_dims=[["paras"]],# R/ o6 }$ K7 f* ~
dask="parallelized",
' [ Z, Q9 A/ n [* b output_dtypes=[np.float],& o& a* y [; v: Z& v% B/ O0 o
dask_gufunc_kwargs={output_sizes:{"paras":2}}2 o0 w D9 T. ^( H4 x- e: J
)
9 a3 n0 D! F( N* K3 w, s( @+ u; O" I5 X; F& i5 `+ L
#compute进行真实的计算并显示进度
. ~8 ]0 {6 w- l+ |+ ]& @/ a with ProgressBar():$ c' V* v8 p8 i. F
result = results.compute() A# K* O1 n2 i1 w. ?2 ?/ x, D' T
* J A, E! M) h T2 | #结果冲命名保存到nc文件
7 q8 c' k) x& K$ H/ s; |6 `' a result = result.rename("norm_paras")
+ j4 v! @2 J5 I result = result.to_netcdf("norm_fit_paras.nc",compute=False)4 }# g: t0 w/ t+ R
with ProgressBar():) A- L6 e3 i) n n& O8 Q& ~+ Q
result.compute()
R! n, o# F/ G3 Y* Z 转自：气海同途 ) D! q; n5 `' F6 K6 U q8 g
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
0 g! O& ~3 C; c2 b' M s) S) l* b% f1 r; u G: P
推荐： 9 Y7 M+ k; M+ O! x0 R% Z' Z! S5 v
1、Python语言在地球科学领域中的应用实践技术应用 ! `5 c. q% s) _+ U% [: B
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
9 Y+ n! X G. A9 ?4 e, P 3、Python在气象与海洋中的实践技术应用精品课程 9 w& h5 X( b1 K$ N3 d& g
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 * `/ |; P4 d/ D, D
5、全套Python机器学习核心技术与案例分析实践应用视频 + J/ H. w5 ~0 A1 x2 z) G# a! `7 s
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
* q6 D- {0 V. W$ a 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
3 Q7 s& s# q/ K( o$ H7 ]" x6 Z# f9 ~. g$ H, n
( f! f6 _ i% E2 ] p+ g7 n
3 E8 X8 |8 n( t
5 R- p" O) i5 @& B" a( c7 c | 
