|
" X, o3 A$ |9 L. s0 `5 ]6 R, {
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 5 a$ H- |2 v6 l( K+ ~
/ p0 g6 D( m3 A& v( D3 c# R1 @" B 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 9 s, M# e% {: w: E$ x! T
Q! C" z% f( U2 j' L3 n 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
% r8 O' t, ]' z0 U' w4 M# j
& ]9 _' A9 M# O4 n 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 ! j2 M& K6 `) Z1 U$ v. a
其中几个关键点解释一下：
! ^# `9 W* Z* N. @# N( ?" o （1）首先定义拟合正太分布的函数
2 W2 I2 ^" X q5 D) z: G def norm_fit(data):
$ N ~0 ~: f) h; e7 s loc, scale = st.norm.fit(data)( p. J1 t. m9 ~" P' S2 G
return np.array([loc, scale]) 8 |2 U' J. ~" B) B" Q) j
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
4 N2 U" T8 A, s （2）xarray分块
) O) _, x, N( p7 I x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
0 C) u( g L7 i2 Z/ ~3 m1 I# Y+ [3 p2 W x = x.chunk({"lev":5, "lat":100,"lon":100})
6 v. e' T* n9 u( n xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
/ }* a7 Q. d4 h) c2 ] （3）xarray.apply_ufunc函数
3 z; }- Z" I( m. u8 X result = xr.apply_ufunc(norm_fit,; @$ s, O" J& G+ W; @! M, }/ K
x,
/ g1 {3 Y V9 B: M7 H. g; k input_core_dims=[["time"]],4 _$ w& X3 i0 s4 R" f* }! i, U
output_core_dims=[["paras"]],
4 H& l9 E4 Z# Z: [" A; W/ O dask="parallelized"," Q9 u# x/ J9 N, |6 D0 u" j
output_dtypes=[np.float],/ Z, S, p! z+ g, u0 u
dask_gufunc_kwargs={output_sizes:{"paras":2}}) w' f% ~, U( g7 B# s( B
) ! g& ]0 ]$ r: ?& V6 T5 p$ c8 `+ ^
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 0 I3 H, j4 w% l3 s; [
9 R4 B* z8 p7 y6 Z p( e! N: B$ S
以下是全部代码：
9 P) X' U* h1 }! R# r7 D0 r from scipy import stats as st
, E/ m/ f# ~; i: @: B. ]3 h7 F" Z5 [1 a. H: b
import xarray as xr# ]4 A8 j+ t- D2 E1 a9 P
import dask3 t1 j6 s' O' U9 {8 P
import numpy as np
4 S4 i2 }9 {+ b, D2 F u& T- p from dask.diagnostics import ProgressBar" j# w4 I8 N/ B1 j* U
/ ?0 L h/ ~2 ?% e' p) M def norm_fit(data):
( p; @% F$ v! g8 M! F! B( w0 n loc, scale = st.norm.fit(data)# a( q! ~4 E& E. I4 j1 P
return np.array([loc, scale]); X& m/ O8 b: o; x
% {* w t7 }6 U V% R
rs = np.random.RandomState(0)5 f+ h1 _2 O4 C3 \8 `- D
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])3 s3 n A! b' b2 Z$ c
x = x.chunk({"lev":5, "lat":100,"lon":100})+ {% J. o$ }- H- ?' h' I1 q
3 G) c) X- h% q #使用apply_ufunc计算，并用dask的并行计算
) d: W- K8 ^/ O0 q result = xr.apply_ufunc(norm_fit, Q9 C6 N! m) E: u* Q, {7 U, ^
x,8 F& t/ i3 I. a: L7 i- _9 U
input_core_dims=[["time"]],! U+ t5 K3 _* D- ^% E- _
output_core_dims=[["paras"]],
& Y7 a4 o( Z- [( p4 d# z, I dask="parallelized",* k4 l/ F- A- k+ ~2 R2 d) H1 e2 u' ~
output_dtypes=[np.float],/ Y7 j3 `- A: |8 r; S/ c2 [' m+ Y' X
dask_gufunc_kwargs={output_sizes:{"paras":2}}
: R% [1 z) J' j )3 e0 L) h I3 T* @% h
: k/ E5 n7 w7 Z9 {6 Z' K
#compute进行真实的计算并显示进度
2 m: M5 I3 K+ [ with ProgressBar():+ O* q5 S3 x( ?/ z
result = results.compute()
3 P; ?2 K$ r& v+ E- \( E& ~4 C5 T) e" P8 V% P' ~0 b' G ?4 \# L
#结果冲命名保存到nc文件: O& r3 q. c6 P& R5 K2 A
result = result.rename("norm_paras")4 c9 I( u% Q* U3 e
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
9 b; P' s6 _# m9 h! N with ProgressBar():
- D, l3 s- l0 [ result.compute() 8 i1 ], u9 Z% d6 @+ o# l
转自：气海同途 . w3 \0 P, f8 [5 p) Y2 P
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 # d+ O6 d! ?" J) p
; k5 G0 E+ W/ D* K- R4 X* z$ u, K 推荐： 1 v; H( C/ h/ Y" v& T& A' K
1、Python语言在地球科学领域中的应用实践技术应用 / O! F5 ^) d3 Z5 i1 ?* O7 r: X% g
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
% c# p0 d0 `( D$ w+ d- D/ } 3、Python在气象与海洋中的实践技术应用精品课程 1 S5 C+ e# _% n r6 ^
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 & S. G6 u0 `. h) P; Q( W$ u! Q4 E
5、全套Python机器学习核心技术与案例分析实践应用视频
& x% y m2 q3 R% ~& b 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
, ?' ]) y$ y! B0 F; q" a 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 ; J4 l# S; b8 I x5 ^# N
2 d" Q" F$ ^8 ^8 t: W& u( {( ~6 z( P, e. P0 A
% Q2 W: l0 g$ G0 q, `
; |. _% [% q0 s" @3 [& J. [0 y | 
