|
& N- O: h, |; [ { Q 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
7 U# ^) I; J8 t; \/ L6 h% i 4 H0 N, \! X) H0 F9 \
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 ! V6 `9 R& R( `8 @" M4 {4 j* m, Y
9 i; p0 ~2 q0 b, h4 V; {5 J 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
# A% G1 u3 R# e3 K/ c4 L
; m: T! i; \$ G 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
. _3 |2 z$ ?" y5 s6 e% ~/ ] 其中几个关键点解释一下： 4 @- p6 e- d/ C/ B# J
（1）首先定义拟合正太分布的函数 % [& x0 g7 X8 [- F P- U+ M
def norm_fit(data):5 h# c+ y: J8 L- m, E
loc, scale = st.norm.fit(data)! b* u; ^- y2 h$ Z! X
return np.array([loc, scale])
B. z. @, s3 |! J 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
# g% \% {* q0 s# H （2）xarray分块
# s% q+ {8 k4 w V x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])" h9 t4 m. ] e( f+ V4 w4 y
x = x.chunk({"lev":5, "lat":100,"lon":100}) ; d4 `) a: i( j* m4 X0 N
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
; ], J( R8 B% c5 g- Q* [ （3）xarray.apply_ufunc函数
' ~: _/ z ^6 t result = xr.apply_ufunc(norm_fit,
, M9 v& }7 Z% v$ A2 E/ Z! a x,: X; N ?, X' O2 P- C# \
input_core_dims=[["time"]],
* d& ~1 H7 I1 h9 K, v1 Q( V output_core_dims=[["paras"]],
. r$ B! ?6 \/ ^' ?+ a4 i8 ] dask="parallelized",
3 z' E9 e* Z8 l$ r output_dtypes=[np.float],/ \7 P0 Z( v. P9 |! i; r
dask_gufunc_kwargs={output_sizes:{"paras":2}}
& |6 E' i. {8 |8 @' r: ? ) . G0 _. v/ y. T& P9 q3 J6 T
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
" r: `# n+ R3 }
6 ], [7 g, r$ q; A/ a! o# N$ r 以下是全部代码：
5 }, u- C+ z2 |3 |; Q/ i( w& u from scipy import stats as st
" a8 t: m" ~/ t; q5 [5 S( ^$ y
( P" C `* B' x- r3 A* G6 M import xarray as xr
6 }/ t/ I+ @! { | import dask1 Z& p I7 Q$ Q d
import numpy as np* v; n0 e" T4 T# E u5 X9 {
from dask.diagnostics import ProgressBar B2 d6 Y6 h; ^5 V" S' }
, i4 p3 c6 l( U/ |. g def norm_fit(data):
9 g( Z: A) ], C loc, scale = st.norm.fit(data)1 R3 G# n. r7 h% [, N$ n7 g
return np.array([loc, scale])! e8 q- j+ s' W* d" d, }
! p- g# B$ Z7 [! N
rs = np.random.RandomState(0)7 r/ \/ m" Z" f3 t; D
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]); _% l2 Q# }2 a, T7 f
x = x.chunk({"lev":5, "lat":100,"lon":100})
s* z8 ?, F) c* X: Q+ I; S
. h* D/ d# L6 f; m #使用apply_ufunc计算，并用dask的并行计算0 X \, P$ V" l8 L& P( g
result = xr.apply_ufunc(norm_fit,- t8 _$ Y, n2 ^4 x# ^) c O
x, a% W- G! l7 J( o+ i1 p: O J, m
input_core_dims=[["time"]],# O- e5 s# s/ r; ]) Y
output_core_dims=[["paras"]],; Z" ^9 _& w# g3 \. o8 r3 I
dask="parallelized",9 e/ c m% O; y& I0 m
output_dtypes=[np.float],
( i1 A; x/ u7 [" V dask_gufunc_kwargs={output_sizes:{"paras":2}}
) ~+ z( H) n' g9 V2 @, Z: } ): q k L5 F- H- X- H/ m! w
7 k: [, C9 n; G& w: k( t1 k #compute进行真实的计算并显示进度 |9 k* z6 K: x
with ProgressBar():
, | u4 |" i$ ]. J. A3 i result = results.compute()
2 n% ~+ N3 g: E( Z
/ Q) U5 T. l8 D% Z4 D. ?8 X# K0 P #结果冲命名保存到nc文件# Q" J" _3 v9 ]5 i; B
result = result.rename("norm_paras")- E0 _, |% [+ [7 _& ^) d) `& a
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
5 E- F1 }/ C6 k9 [+ B" ]* e9 d with ProgressBar():
( U0 p: J2 F1 I" N$ u result.compute()
3 ]) m+ V5 E& u0 ~) ^, G 转自：气海同途
$ u8 z% D& Q1 J! \ 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
0 k6 F; r, Q. a5 Q3 X. }
7 q" E8 {; l0 j: O7 K 推荐：
* D# z- A% [( R A+ e 1、Python语言在地球科学领域中的应用实践技术应用 3 O9 s' e5 f5 s3 [5 e
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” 2 Y4 @, |2 h* }. S4 |4 x) N J7 j/ K
3、Python在气象与海洋中的实践技术应用精品课程 9 i/ B% V0 F* e/ q8 a- j* X, Q4 C
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 5 n0 c' R9 a; p* b: ]9 r; o* V
5、全套Python机器学习核心技术与案例分析实践应用视频
; d* u" C" r, T% j& E 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 ' O7 A/ ?4 ~3 O
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 4 Q) C7 s1 O2 q/ ?
* g z$ E: o0 ]0 B
, ^9 x6 ?7 \/ L8 `, H9 C
( {9 E& U# b6 K3 \' X
4 C/ I" p# N _3 @6 w
| 
