|
/ {2 p2 {' T' _
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
# [" d* d$ Q7 |. o7 k' H 0 }' \* h! s' X. O
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
/ V4 O8 k8 x! e7 i- R9 E! r# r+ w: l
. q, x5 F. @1 J 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 " ]- W( ?' t& p+ Z2 |
7 d. U7 a) W& M) S( }
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
0 I- V) K' O; W" [# Z 其中几个关键点解释一下： 0 F7 i, d3 p+ Q) |7 V" e' }5 ^: F/ A
（1）首先定义拟合正太分布的函数
& I3 Z/ M. h' H. y% I3 l def norm_fit(data):
& }0 N' o" r3 z5 h1 J loc, scale = st.norm.fit(data)
3 w/ U) T: s: W3 I return np.array([loc, scale]) # l- Z6 g, i) ?- h4 E1 Q
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 " Z' J+ ^. V* I2 ?' E
（2）xarray分块
" k' F/ H# S3 C1 o9 | x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
- d. R6 _0 W+ N4 B4 h- l% W x = x.chunk({"lev":5, "lat":100,"lon":100})
) C6 B$ [' D/ _; I- _ xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
& {7 W4 {! v: v0 i （3）xarray.apply_ufunc函数
9 W- ^5 L# a5 W$ O6 O; w! ~3 G result = xr.apply_ufunc(norm_fit,
" U, P$ O: q0 r2 s9 @, w+ W9 _ x,& l& v4 u% [6 f
input_core_dims=[["time"]],
6 o, D/ K& K& R3 n# j5 y output_core_dims=[["paras"]],' [; T: w5 O) I0 \& `$ E) K
dask="parallelized",, x* y: b5 C! E. O% |
output_dtypes=[np.float],
2 H1 x) J1 L( _, V! c) a dask_gufunc_kwargs={output_sizes:{"paras":2}}' K6 g6 d5 i; n, w9 ~
)
3 \" Y! z- b8 c( O, { apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
A& f' V: C4 G9 U* ^6 g8 K* q
Y- f6 U" I" K2 t9 v! v$ } 以下是全部代码： % C p( F$ o! O5 v! |) @0 g
from scipy import stats as st; l, z# j' `( M# d$ _
2 C. a. ^- v) y) K9 N# }
import xarray as xr( ? L' b, j. Z* r/ n0 l' u
import dask4 F( ^/ c7 t I5 n/ \4 f4 T
import numpy as np. m. s, Q; a6 A$ c
from dask.diagnostics import ProgressBar
, m( y) X% i4 ^: h( T
* D# d$ L4 b6 o8 ?# x n' T def norm_fit(data):; `/ `; z. m S' I& V/ A
loc, scale = st.norm.fit(data)
) i5 t) Q5 w- _ return np.array([loc, scale])6 ~) y2 N5 U/ ]( {, O# C) Z ^2 O) t8 M0 _
3 a& I( g3 t1 P( Y. B. L3 x rs = np.random.RandomState(0)
! z: [# p! a9 ] x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]); N0 f* }1 l$ A* x" ^+ p
x = x.chunk({"lev":5, "lat":100,"lon":100})3 o1 ^& Y( O# V6 w9 [
% L# G8 o& _0 E9 H
#使用apply_ufunc计算，并用dask的并行计算
! ]1 [9 d) q4 v; c! l m result = xr.apply_ufunc(norm_fit,
: q. ^8 M1 y7 j* C3 G. U B! i3 k x,8 u) O$ y( I* j2 q
input_core_dims=[["time"]],
) O1 o! e- Z, ~0 k" A9 U9 L# D' G output_core_dims=[["paras"]],6 U. l" _- e& d2 i$ o- ~
dask="parallelized",
4 |. ^ i% Z# b output_dtypes=[np.float],6 R* f! F! ] A4 c: X$ c6 A3 ^
dask_gufunc_kwargs={output_sizes:{"paras":2}}
, }* m/ O3 g- O& ~ )0 y. w8 k+ v0 F' `
. x* y9 E8 u3 G+ a5 K #compute进行真实的计算并显示进度
9 V+ C& [! O5 }4 C4 S% I" O with ProgressBar():
' ]6 g$ O' c$ | e result = results.compute()
6 g Y, J+ @) u9 J; z% x! S- x, V) ^1 u7 V
#结果冲命名保存到nc文件
5 \2 @# u: e' Z; t result = result.rename("norm_paras")
P) l3 A5 i! g( N, Y; ^/ d r result = result.to_netcdf("norm_fit_paras.nc",compute=False)
: I4 B8 B/ J( `) C: z3 T with ProgressBar():
6 T) e0 }9 U2 r* y( O result.compute() L/ _7 M! D8 u! x
转自：气海同途 / ~8 ^: A" }$ e6 o7 V# B: g
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 7 M2 C% B; G0 s6 H# x$ ^9 b; z X5 O0 z
1 c5 o* v9 `- T+ B# k% f" r1 t% x
推荐： 7 u! b9 q' }; Z# y8 B
1、Python语言在地球科学领域中的应用实践技术应用 0 o2 y- q* {; h5 z+ J) Q
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
/ Y$ j2 ?4 }: K, E J1 x 3、Python在气象与海洋中的实践技术应用精品课程 3 f0 t+ k6 ^5 A7 q g( L s& w0 D
4、Python在WRF模型自动化运行及前后处理中的实践技术应用
4 d5 C) K0 [. k( q1 Y' c. l' o 5、全套Python机器学习核心技术与案例分析实践应用视频 $ J: A( A9 [- c& k3 g- o
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 & e3 _2 i& g O* m+ p; X7 f
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 $ j* G6 H6 d% ^+ S
2 z5 s' T3 I" f0 Z8 H* o) r& B% m; q6 M M. a. b9 J% p
/ Q+ I4 ?3 f% ^
% y+ a* P/ B9 w& I | 
