|
9 D& X& ?6 ~2 _1 ?5 p0 w; G
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 & V J+ G5 }2 C% m7 @! J) T0 w# y
! F/ W3 X4 j# z 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
; U% X* W& d) c; o* V; h! \8 K, Y
5 M- I% H+ l; c* _* Y 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
/ R( Z; {/ o4 L8 z- c1 I5 R% ]! l: I3 _3 C2 ]% B6 b5 f6 c( o4 A
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 ! k7 D; F& _6 l, F. T) |, `
其中几个关键点解释一下： 8 a! O+ \4 v: s/ f! G A6 S% C7 ^$ N
（1）首先定义拟合正太分布的函数 n2 L6 L% ~3 A3 Q3 F/ v
def norm_fit(data):
2 M$ Q2 H' m' y1 P loc, scale = st.norm.fit(data)
W4 M, N. e4 ]& e return np.array([loc, scale])
# C) a. {& j- @9 M4 M2 V* ] 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 6 v3 \0 t4 U( s0 X+ {0 j
（2）xarray分块 * @: C, B% D0 ]# D& K$ ^
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])# P, ?0 x$ A1 p( f: N$ c1 _8 \
x = x.chunk({"lev":5, "lat":100,"lon":100}) 8 P. P# m4 ?' s! N
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 6 L( ~" Y! g H
（3）xarray.apply_ufunc函数 / a+ G6 P( t5 u& A Z0 [4 Z' ^# V
result = xr.apply_ufunc(norm_fit," _: w. ^' N( \6 S( |
x,/ x6 t+ x0 g7 T8 ~
input_core_dims=[["time"]],) j) N7 O4 l4 q3 |
output_core_dims=[["paras"]],/ x* n* }' b7 ]+ D# _# G
dask="parallelized",
4 E- f) i5 p( s: t% n3 _% I& @# w output_dtypes=[np.float],9 B R; z8 |! B# r( M6 [9 X
dask_gufunc_kwargs={output_sizes:{"paras":2}}
* ~* J& [; l; i I( c1 ~ )
% }$ m! x5 d: s/ Y1 L* N apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
% H/ Z8 r# T1 R0 O! y% a* ]/ N& S% m, J
以下是全部代码：
! w0 J' D# q) t: _8 R from scipy import stats as st- v/ I3 i4 {. d$ G: v1 z& x# l
+ N5 U0 o, _9 {6 t% j) v import xarray as xr
' r% Q5 Q& |: s: z' c7 c import dask" h! |( g& z; s# E q. G2 k5 t
import numpy as np
& g$ a6 v% u& z7 n from dask.diagnostics import ProgressBar
) [+ q' U2 B% S* H' Y% U9 V: b4 G% S4 d
def norm_fit(data):
/ s! a' U4 }! Y2 ?) A$ N% a2 h8 Z loc, scale = st.norm.fit(data)
' p0 V8 g0 ]+ f/ ^* h( S. T& J return np.array([loc, scale])9 h" C! G5 Z! r* v1 U9 \6 a, j4 i; s
2 ^' e& B: b2 I- i3 [4 a
rs = np.random.RandomState(0)
0 c- M8 ]: _- {" n0 r x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])& j1 N2 m2 j$ E. a
x = x.chunk({"lev":5, "lat":100,"lon":100})
9 `! ?# F& s U( ]8 R( L6 X% g) s5 ?: z& d9 o: R
#使用apply_ufunc计算，并用dask的并行计算
5 Y* C9 o% \ F4 R- |" S" R result = xr.apply_ufunc(norm_fit,
. I( k6 m" \2 c x,' X% ^# a2 z3 D3 @; s/ A1 ^
input_core_dims=[["time"]],6 y! g0 \5 Q% M7 N: ?$ r
output_core_dims=[["paras"]],
# I" S& q5 L6 a# c1 w dask="parallelized",+ N; X) }, o7 g3 @, a
output_dtypes=[np.float],3 M$ K% C8 A/ E% q9 `; u7 i( i
dask_gufunc_kwargs={output_sizes:{"paras":2}}
8 c4 l8 ~5 |2 U5 @. W' X )
4 {7 Z5 G+ m% j; j4 p5 Y' L2 }( G6 s5 m
- l/ k% v0 X4 n. K( F( R1 Z #compute进行真实的计算并显示进度8 g, a* Y0 K$ C. v
with ProgressBar():
8 \) |! P9 I) B0 X; E result = results.compute()
, B9 [1 ]& Y: o+ N3 r I" l
4 D( B2 D+ M0 t2 n8 c #结果冲命名保存到nc文件9 K1 H1 s9 A, x5 J1 }* i0 h/ z
result = result.rename("norm_paras")* C& W3 h; D. y* i+ s. @6 U
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
" w4 I. O. d" w; p' P$ B% t$ s8 { with ProgressBar():
2 |+ x& B% C0 _5 x' e2 v* G" N* k5 g result.compute() # I- D0 W7 D4 Q. X9 c' X& c
转自：气海同途 $ l. A) m& N1 ?6 e6 i9 G4 [4 a
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 5 H# X' F" E4 a8 X2 y; }! x7 \, e3 w
/ d+ ]" N1 i' S" S2 }% m
推荐： ( h w% R- ?3 |. y% Q1 ]# ?2 L( k
1、Python语言在地球科学领域中的应用实践技术应用
5 M9 d" Z! ~7 y" J& W, E; @0 G$ z 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” . h" G7 G5 O, x, u Q
3、Python在气象与海洋中的实践技术应用精品课程 * J0 p O1 N" ` n. T0 _) S
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 + z% z& |# ^) o
5、全套Python机器学习核心技术与案例分析实践应用视频
' N q+ N6 g+ ~; S! C 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 ! u) ]; G1 u/ O% h* J
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 6 S4 Z/ Z$ r# E$ ]
$ n0 z* l5 A6 z, n" _2 u# q! \
! q( z a- r( _: v' |* L ^, q* p7 ?, L* {2 l) U* X
1 s% n0 o$ `( @9 W! M. u' ?* M | 
