1 y" x5 {6 r, q, G/ f
气象海洋领域,常常会涉及到大规模的数据处理,比如高时空分辨率的模式数据、雷达数据和卫星数据等。 , B$ v9 p# J1 r: }5 D- Y/ b& O
. _5 n7 A' K2 B) d& u' p
大部分情况,我们一般只会涉及小规模数据处理,对计算效率并不会太过追求。但是当数据量变大时,低效的数据处理所耗费的时间非常明显,因此高效的数据处理方式尤为重要。
8 f4 T: K$ d. x0 u! w3 O% a" s8 ^, u3 x0 Y& G! z3 D8 F3 v
本篇以拟合一个高维数据的正太分布参数为例,介绍如何使用xarray+dask加速数据处理。
# u; v" n; \ D% R: K' L& ]1 z9 O/ d: m1 t/ u& }
数据维度X[time, lev, lat, lon],需要对三维空间每一点,沿着时间维度做正太分布拟合(正太分布拟合只是作为例子,这里可以定义你需要的操作函数)。 # d0 f1 g$ S/ o1 P) e# r
其中几个关键点解释一下:
" o& {2 K1 r% @( b' y (1)首先定义拟合正太分布的函数 9 r+ w( E8 h% M2 H1 X" q
def norm_fit(data):# D/ @ J: N6 ^# U2 { [
loc, scale = st.norm.fit(data)7 u+ J) m1 |8 L1 f
return np.array([loc, scale]) * M3 ?4 p) o6 N% I; I4 p
这里需要注意的是,拟合的函数,输出参数,需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
- z! v8 P6 S- u3 o0 Y. }# D (2)xarray分块
4 o5 _6 r( `+ z- ^4 n x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])* E6 H% t. v1 Z% p# L3 P
x = x.chunk({"lev":5, "lat":100,"lon":100}) 6 z! M9 Q0 j# N5 l# T9 B7 ~
xarray结合dask可以将一个大型数组分成一个个数据块(chunk),需要注意的是我们需要沿着时间维操作,拟合需要整个时间维度的数据,因此时间维time不能分块,只能对其他维度分块。 " a& f2 W: x9 q; D# ^' A! w
(3)xarray.apply_ufunc函数
: ^7 O+ e! y1 c6 J6 n" n result = xr.apply_ufunc(norm_fit,) |6 x8 l6 I) W: M# m, s+ N
x,) G: @- `, O- @, I3 m8 S$ d
input_core_dims=[["time"]],
5 t7 ]2 r: j# X2 {0 d/ E* G& m output_core_dims=[["paras"]],
& e( X1 M# l2 h) ]5 R+ u4 d dask="parallelized",
" U( B% I0 i% [0 x output_dtypes=[np.float],
r/ }9 S% A( v5 u dask_gufunc_kwargs={output_sizes:{"paras":2}}) @ k3 X' V9 \* a# w: \7 C3 B
)
! i9 E( t: e2 j& b apply_ufunc函数具体可以参考官网教程,这里只说使用时遇到的困难,即如何定义输出维度:输出维度是用output_core_dim定义的,将输出的拟合参数(期望和标准差)定义为paras维度,维度的大小为2,通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据,在每个空间点对时间维度拟合之后,输出的数据为[lev, lat, lon, paras]。(PS: 这里感谢深雨露大佬的指点)
7 ]1 u% U) \# Q! ~
! ?3 P0 M, }% t' L 以下是全部代码:
- r; j2 q0 u! l4 Z! `3 N% q from scipy import stats as st4 n8 Q( g5 k4 [/ }8 A: v$ L* T; W
8 c8 m7 F: r% ? import xarray as xr6 R! e' |% L/ i
import dask3 L7 }5 F; ?* m! U9 w' c- h; q
import numpy as np
# L( \$ c# l5 k4 T from dask.diagnostics import ProgressBar
. z) n$ N( X, r: \6 Z% w1 e2 Z, |+ |2 R+ T* V
def norm_fit(data):8 }% K0 N/ D$ o2 N& ~* ?8 [
loc, scale = st.norm.fit(data)
& u" G/ s- P$ h return np.array([loc, scale])
: g. t0 U2 y. z+ C9 X4 d) H* `; t' n& x" ]$ u2 H0 z
rs = np.random.RandomState(0)9 p: r {! \" \+ G* I6 w
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])4 i! \( y3 o0 L( K) J* c8 X
x = x.chunk({"lev":5, "lat":100,"lon":100})
* h0 D' o# H$ W# ], j! }2 ~# J( E( M$ x& s5 u9 I+ {$ h: A
#使用apply_ufunc计算,并用dask的并行计算
& x2 E3 Q) n0 T# `+ R result = xr.apply_ufunc(norm_fit,) ^! K% V+ g6 V2 W: y& S! v3 f
x,0 K3 j l- ]2 A! s; O# _
input_core_dims=[["time"]],8 m* F$ ~1 N j: w
output_core_dims=[["paras"]],
2 P( d7 G- U( t4 D5 W dask="parallelized",7 X' v" S1 n* i' [6 ?
output_dtypes=[np.float],0 C1 x+ I y2 M' I0 d" Y( a
dask_gufunc_kwargs={output_sizes:{"paras":2}}& I0 B6 n+ @# A+ y$ k, E
)* q2 c' F" b! \1 @
4 f0 f6 Q S- b* |& `* K; b- v #compute进行真实的计算并显示进度+ K9 G# s \9 I! N/ i
with ProgressBar():
) u3 r% A3 |6 P8 f3 L- T @2 T result = results.compute()9 `' s, ? a1 W( I) y0 T. D: R
8 {1 I7 V: c7 w! W) e6 n6 a
#结果冲命名保存到nc文件; x6 c. |+ @0 k1 n
result = result.rename("norm_paras")( K6 C2 M6 F$ d! \3 W. M+ b9 a
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
+ X( G' M& x6 p. S; n) _) A9 G with ProgressBar():
: P/ h7 v$ s5 m- E result.compute() ' U+ W. X. `( A/ v
转自:气海同途 1 U; b6 Y$ L' H+ s- |2 ?
关注【Ai尚研修科研技术平台】公众号,查看更多课程安排及免费数据资料
8 P( N7 R% ~% w4 ~' z j5 X. [ / z& r7 f' Y* N" u8 W
推荐:
4 u, T. t7 Z4 f4 Y7 N# L 1、Python语言在地球科学领域中的应用实践技术应用
, `" H7 n- ?6 M6 I ]2 A5 F* S A 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” + }3 ^9 w: l6 }" \1 f5 M. a+ E
3、Python在气象与海洋中的实践技术应用精品课程 $ x0 x9 I+ v5 n
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 7 ?- l' S" H. G* l( h" a. @
5、全套Python机器学习核心技术与案例分析实践应用视频
! a" C7 k- b5 Y- d0 E3 H+ U7 w4 X) r 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 * |. H, p1 ^( V7 I# _% Q3 L
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
2 }4 m/ v, M4 x- g) y8 r% m
9 r5 l l F" ^' A7 I! n
# B M+ s1 r/ H" j3 w% t) {; g0 h4 T# y4 t/ T3 x8 W( X7 [" J
( v. @' \5 \# l' r4 O0 @& W
|