|
- B+ @+ W4 }& i' n! e( s
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 : C5 q# g* k$ h! g$ i# R
k) K9 c* Q/ `, x0 p 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 , o9 W/ f6 F$ x
0 R( V3 ^1 V. w6 z2 y 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 $ S9 n7 V4 a& l- b8 R' n+ D0 H
5 F& e3 |+ Z7 l) A
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
$ e3 R4 V6 Q$ v9 Y# n# T$ ? 其中几个关键点解释一下： . ^. h& w4 t' u
（1）首先定义拟合正太分布的函数 / K& K/ ~* L' @1 {
def norm_fit(data):
5 \4 ], O( e3 y7 i3 s loc, scale = st.norm.fit(data)- ?4 S0 M+ ]2 i2 U. ?/ ?
return np.array([loc, scale]) , b/ r2 [2 k- S! q
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
6 X# H) `6 ?9 C （2）xarray分块 / M8 w, y7 b& I5 n- a/ e; i9 E
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
1 q+ H% _' P& \& F x = x.chunk({"lev":5, "lat":100,"lon":100})
1 t2 N+ G& O3 R2 l. x6 a xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
+ p& n4 O9 ]2 {. x E4 s （3）xarray.apply_ufunc函数
6 v( Q& v& \1 H3 d+ i1 |& } result = xr.apply_ufunc(norm_fit,$ f* F7 U. i* ~
x,
, S& P& ^. K' H5 X' t input_core_dims=[["time"]],, l0 ~+ T }2 A, |, s: H
output_core_dims=[["paras"]],
; |* Q, x( h: y! h$ J' _ dask="parallelized", W2 U I- x$ r$ J2 `: S
output_dtypes=[np.float],, d6 c+ r$ E0 C5 d/ `4 {
dask_gufunc_kwargs={output_sizes:{"paras":2}}: Z. ^/ k1 A `# j6 V, S6 S8 j
)
4 @- h' u5 J- u% n _ apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
1 P& A" Y' H8 L9 t7 Z0 d5 N3 Q# D% p3 [; Z5 N8 C
以下是全部代码： 6 c5 E8 y* A3 n4 E6 ]* x
from scipy import stats as st
( I. Z6 j- E2 x, q% s
& S2 a1 c/ E! `* R% ^: [ import xarray as xr6 b4 c7 @- C3 W/ K
import dask
- b" B6 @! ?' s+ `! ^2 I/ q import numpy as np
% o- p( ]) m; u! N3 r F& w6 t* l; T from dask.diagnostics import ProgressBar
) ~8 `/ G: R: |2 l
% \( |8 Z5 V' L) ]7 ^ def norm_fit(data):
: z, c5 E# E8 O% y5 E loc, scale = st.norm.fit(data)
- g+ \# o. [1 U9 t return np.array([loc, scale]). t2 E. w' L. U' _& L
6 p* k; K4 U/ b- ^0 p) f
rs = np.random.RandomState(0)$ j% Y" Y0 i6 D0 K$ h/ z" m
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
* E4 g* Y5 m2 m3 z: N3 u" o x = x.chunk({"lev":5, "lat":100,"lon":100})" B$ g. _# H- Q" I/ l- Y! d7 X
& R) `0 S5 r4 {& k2 j3 {# ]+ s" u #使用apply_ufunc计算，并用dask的并行计算
0 y8 v2 |9 ?8 f- Y. [+ a- ` result = xr.apply_ufunc(norm_fit,
& I ^8 o# d% P( w$ r x,
, M, f `) w$ p4 s# \ input_core_dims=[["time"]],( T' ^$ Q. r/ J/ h! e
output_core_dims=[["paras"]],
( }& d: J: e9 K4 N2 V2 C) P dask="parallelized",
# I6 M7 A, F) ]& l; R" E0 a output_dtypes=[np.float],
& K8 [' P! K7 F" c" m2 ^: ?+ O1 n G dask_gufunc_kwargs={output_sizes:{"paras":2}}
3 U4 {7 H+ b5 j4 g+ G$ V3 P1 [9 S )
% g$ _4 C4 ~- r4 j }/ M; L+ }& Q6 b3 d3 Y1 V
#compute进行真实的计算并显示进度) f- Q- z0 t& {, B7 o2 K
with ProgressBar():5 u0 a+ d$ M. X ]$ F" t- y+ `' ]$ ^
result = results.compute()
7 l1 a0 H. ^3 m( p
7 D- s$ ^! M; Y4 y, E5 K #结果冲命名保存到nc文件2 F: z0 P) {; @" Z v
result = result.rename("norm_paras")& ~9 s- F( {7 M% I. j
result = result.to_netcdf("norm_fit_paras.nc",compute=False)( ?1 P' d3 k* s# d9 a
with ProgressBar():
! Q5 [' _% c* x( g result.compute()
" R% m$ r* Y, J& `* W/ D 转自：气海同途
4 \6 F! x0 g5 @ 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 " l! p! L2 Q$ y& F1 |
- s2 ] |+ ~; h* u2 v* Y 推荐：
) _; F8 a$ Z+ k* \5 H" ^. p 1、Python语言在地球科学领域中的应用实践技术应用 & x; v2 _: g' C0 n7 O$ [ W
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
7 X9 {# y( U" [* y 3、Python在气象与海洋中的实践技术应用精品课程
3 Z2 i( S! d9 U- }& D- ^$ g1 E 4、Python在WRF模型自动化运行及前后处理中的实践技术应用
+ x2 l9 |7 `6 }, I( m1 O8 f7 a/ s 5、全套Python机器学习核心技术与案例分析实践应用视频 + V0 Z7 z. p7 c! X( ~) _
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
5 Q& t- P+ H5 B2 }3 h) [/ V& D/ G7 \ 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
$ F T3 ~4 x8 q X, @0 S8 R% \& Y o7 ` Z: A* W
4 _. ~- J; M: t1 `: i5 M
8 O. R' ?8 u5 m: ^% @
4 d+ u1 L; E" u# H | 
