|
9 ~. k9 ~5 Z7 c( @7 z. j, S 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 9 T7 \& p* {, H0 B
9 r! N# Z2 r% w( B8 q2 j
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
. p* V$ y6 e! Y R- J' P5 b: x1 x
2 h, U& }7 M- O$ k+ t. P5 @ 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 * r o' h+ @. Q' A# C k1 }
% K) d8 N3 o3 V# n9 ^* Z 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
& D. w# X9 x. m! u( Y: n% P$ Y 其中几个关键点解释一下： : G1 a0 f9 }* W
（1）首先定义拟合正太分布的函数 ) ^7 o# ?' P. h7 S
def norm_fit(data):
2 I4 l0 q/ }/ o8 z6 k3 g+ x, l# ?- y loc, scale = st.norm.fit(data)- `" E5 X* p2 I _, {
return np.array([loc, scale])
. S$ L0 ^: L9 t& ^( E 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
9 B" K N3 `" P （2）xarray分块 5 a' c- B/ L- q7 t2 }: h" R& M
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])! ~/ w* d0 Y3 P! y# p3 M0 M
x = x.chunk({"lev":5, "lat":100,"lon":100})
9 B! m: ^6 n) B- ^ xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
6 ~/ {* v3 S; m2 @ （3）xarray.apply_ufunc函数 0 K, j0 x1 v- b0 |3 x) z/ x' L; n
result = xr.apply_ufunc(norm_fit,8 r+ {- n/ ?. X- f$ M
x,8 D2 ^' T7 R- I
input_core_dims=[["time"]],8 P- m: o+ ~& Z$ m
output_core_dims=[["paras"]],
+ \: n& O0 }4 H. Q+ f3 o. ~ dask="parallelized",
/ J( d' u2 s( C; F1 D/ V output_dtypes=[np.float],
, U! e* o& Q4 T; i9 ]/ E+ Z; V dask_gufunc_kwargs={output_sizes:{"paras":2}}% X* _7 Y5 ]; }# E; Q
)
- R1 U& `8 h! F Y6 |) T apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
# ^9 s0 d9 U" ~0 C) U4 H
% p5 R: @5 a0 t 以下是全部代码：
4 s5 b! l- {4 W( e4 [ from scipy import stats as st
# x* S- |5 l$ h% a+ k/ c0 V
5 ]. W7 ]2 f4 | import xarray as xr1 h. W" K" `# e* m7 Z
import dask
3 i) G' F i& K1 ` @1 \7 k import numpy as np H9 D$ ]4 Z1 J9 p! L; n$ G6 R' U& _
from dask.diagnostics import ProgressBar0 _2 m5 Z& w \/ E1 ^1 s9 }) `4 S- X" L
( E; [) q9 e+ a4 N) g
def norm_fit(data):
; f" C: w2 {4 o2 m: z$ R$ `9 A loc, scale = st.norm.fit(data)& P% W9 `6 A* ?1 r( b. x
return np.array([loc, scale])
! \: p6 j3 N/ {6 I% O: B8 t+ p
3 [) a# f4 U2 Y rs = np.random.RandomState(0)' j! ^- Q0 d, Z, l; z. q( \ p' A; D
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
. V" @1 }8 T3 X. j, c6 T x = x.chunk({"lev":5, "lat":100,"lon":100})
$ E' Z/ v4 [( X' [( _4 Q6 \" r! ?; w- m- r( i
#使用apply_ufunc计算，并用dask的并行计算+ s1 a. o! {3 L! @' l
result = xr.apply_ufunc(norm_fit,% B# \5 S; H$ Y% N8 v$ n
x,3 g1 L+ L$ B/ z4 n: A1 l% D; J$ C
input_core_dims=[["time"]],: W. W$ q6 K+ ~) d1 A8 I
output_core_dims=[["paras"]],3 ~: N/ s. @# m" d9 n
dask="parallelized",
4 E! x5 I# \( A output_dtypes=[np.float],
9 n8 n- i2 q+ _ dask_gufunc_kwargs={output_sizes:{"paras":2}}
. N/ t: W+ P+ {2 s )) u- y: A1 i- V2 n$ u/ M
# o5 m0 K$ r* a3 p' t# p
#compute进行真实的计算并显示进度
6 h# V* h! {& z4 i with ProgressBar():6 D. W. |9 \ ^' g6 ~7 ]
result = results.compute()
+ @& V) ^$ }2 N8 {( v" f6 R. }+ q
# F' W" W6 N5 h2 M #结果冲命名保存到nc文件& O( K9 T# k- L
result = result.rename("norm_paras")% M h3 r# y9 W; V: p. ^* L: X
result = result.to_netcdf("norm_fit_paras.nc",compute=False)7 [1 u/ }6 f$ j4 z: \# R
with ProgressBar():# d ~0 Y1 Y/ ?! p: L9 M1 c
result.compute()
; U3 ^- _ h+ h9 Q: G 转自：气海同途 ! }+ X2 l% C: j+ l" g
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 & _# W1 C+ }, b- z! `9 {- M
7 P8 H8 N8 K) ^6 {8 E. m0 t
推荐：
3 |* q9 |7 R9 B, A& D. q! f 1、Python语言在地球科学领域中的应用实践技术应用
: F3 _7 o2 Q! Z6 i" t ^- I/ u 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” ( h. v C/ S$ X q
3、Python在气象与海洋中的实践技术应用精品课程 + c" V# V! _ n
4、Python在WRF模型自动化运行及前后处理中的实践技术应用
1 m5 }! n3 e0 K+ } 5、全套Python机器学习核心技术与案例分析实践应用视频 5 C) g2 |7 ^* S
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 # n. k+ w8 `" d; O* A `0 Q
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 8 v$ A& i9 _1 I& [# V$ y4 @5 @
6 n9 Y* B/ ^$ ?8 v' w8 [4 [# P, L% [9 A% I
% P' a) l6 L8 J6 a+ A
9 A- L. ~% O% C. e# O7 x
) V% l2 B4 |/ j* |" B0 S d" F. G | 
